Skip to content

Commit

Permalink
Applied fixes to all files
Browse files Browse the repository at this point in the history
  • Loading branch information
mb2055 committed Aug 30, 2023
1 parent 4756eb1 commit 6eea673
Show file tree
Hide file tree
Showing 3 changed files with 37 additions and 36 deletions.
30 changes: 11 additions & 19 deletions src/somd2/analysis/_analyse.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,6 @@ def __init__(self, data_parquet, custom_meta_key="SOMD2.iot"):
self._dataframe, self._metadata = self.parquet_to_dataframe(
self._parquet_file, self._meta_key
)
print(self._metadata)

@staticmethod
def parquet_to_dataframe(filepath, meta_key="SOMD2.iot"):
Expand Down Expand Up @@ -94,7 +93,6 @@ def _calculate_beta(self):
from sire import u as _u
from sire.units import kelvin as _kelvin

print(self._metadata["temperature"])
self._beta = 1.0 / (
(_const.gas_constant / 1000)
* _u(float(self._metadata["temperature"]) * _kelvin).to("K")
Expand Down Expand Up @@ -274,7 +272,7 @@ def extract_data_TI(dataframe, metadata):
return df

@staticmethod
def extract_data_MBAR(dataframe, metadata):
def extract_data_MBAR(dataframe, metadata, lambda_array=None):
"""
Extract gradients from a processed dataframe,
formats in alchemlyb-compatible dataframe format
Expand All @@ -293,12 +291,13 @@ def extract_data_MBAR(dataframe, metadata):
-------
df : pandas dataframe
Dataframe containing the reduced potential values, formatted for MBAR"""
try:
lambda_array = metadata["lambda_array"]
except KeyError:
raise KeyError(
"No lambda_array found in metadata,unable to perform MBAR calculation"
)
if lambda_array is None:
try:
lambda_array = metadata["lambda_array"]
except KeyError:
raise KeyError(
"No lambda_array found in metadata and no lambda array provided,unable to perform MBAR calculation"
)
import pandas as _pd

temp = dataframe[[str(i) for i in lambda_array]].copy()
Expand All @@ -309,7 +308,7 @@ def extract_data_MBAR(dataframe, metadata):
temp.index = multiindex
return temp

def analyse_all(self):
def analyse_all(self, lam_array=None):
"""
Function to call analyse on all parquet files in the directory
returns a list of dataframes that is sorted and ready to be used
Expand All @@ -318,22 +317,15 @@ def analyse_all(self):
extracted_dict = (
{}
) # dict to store extracted data, needs to be sorted, python 3.7+ required
lam_array = None # Used to check that all lambda arrays match
lam_array_local = lam_array
for parquet_file in self._parquet_files:
temp = analyse_single_lambda(parquet_file, self._custom_meta_key)
analysed = temp.analyse()
meta = temp.get_metadata()
lam_curr = float(meta["lambda"])
if self._method == "MBAR":
if lam_array is None:
lam_array = meta["lambda_array"]
else:
if sorted(meta["lambda_array"]) != sorted(lam_array):
raise ValueError(
"Lambda arrays do not match across all simulations"
)
extracted_dict[lam_curr] = self.extract_data_MBAR(
analysed, meta
analysed, meta, lambda_array=lam_array_local
).dropna()
elif self._method == "TI":
extracted_dict[lam_curr] = self.extract_data_TI(analysed, meta).dropna()
Expand Down
23 changes: 14 additions & 9 deletions src/somd2/runner/_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -192,7 +192,8 @@ def get_CUDA_devices():
@staticmethod
def zero_CUDA_devices(devices):
"""
Set all device numbers relative to the lowest (the device number becomes equal to its index in the list).
Set all device numbers relative to the lowest
(the device number becomes equal to its index in the list).
Returns:
--------
Expand All @@ -208,7 +209,6 @@ def run_simulations(self):
--------
results (list): List of simulation results.
"""
print("here")
import concurrent.futures as _futures

results = []
Expand Down Expand Up @@ -307,19 +307,24 @@ def _run(system, map, lambda_value, lam_minimisation=None):

# set all properties not specific to platform
map = {
"Integrator": "langevin_middle",
"Temperature": temperature * kelvin,
"Pressure": 1.0 * atm,
"integrator": "langevin_middle",
"temperature": temperature * kelvin,
}
# Pressure control. Only set if the system has a periodic space.
if (
self._system.has_property("space")
and self._system.property("space").is_periodic()
):
map["pressure"] = 1.0 * atm
system = self._system.clone()

if self._platform == "CPU":
if lambda_value is not None:
print(
f"Running lambda = {lambda_value} using {self._platform_options['cpu_per_worker']} CPUs"
)
map["Platform"] = self._platform
map["Threads"] = self._platform_options["cpu_per_worker"]
map["platform"] = self._platform
map["threads"] = self._platform_options["cpu_per_worker"]
try:
df = _run(system, map, lambda_value=lambda_value)
except Exception:
Expand All @@ -331,7 +336,7 @@ def _run(system, map, lambda_value, lam_minimisation=None):
self._remove_gpu_from_pool(gpu_num)
if lambda_value is not None:
print(f"Running lambda = {lambda_value} on GPU {gpu_num}")
map["Platform"] = (self._platform,)
map["platform"] = (self._platform,)
map["device"] = (gpu_num,)

try:
Expand All @@ -350,7 +355,7 @@ def _run(system, map, lambda_value, lam_minimisation=None):
df,
metadata={
"lambda": str(lambda_value),
"temperature": str(map["Temperature"].value()),
"temperature": str(map["temperature"].value()),
"lambda_array": self._lambda_values,
},
)
Expand Down
20 changes: 12 additions & 8 deletions src/somd2/runner/_sire_merge_runsim.py
Original file line number Diff line number Diff line change
Expand Up @@ -88,15 +88,19 @@ def _setup_dynamics(self, timestep="2fs", lam_val_min=None):

if self._minimise:
if lam_val_min is None:
self._system = (
self._system.minimisation(lambda_value=self._lambda_val)
.run()
.commit()
)
try:
m = (
self._system.minimisation(lambda_value=self._lambda_val, map=self._map)
)
m.run()
self._system = m.commit()
except:
raise
else:
self._system = (
self._system.minimisation(lambda_value=lam_val_min).run().commit()
)
try:
self._system = (
self._system.minimisation(lambda_value=lam_val_min).run().commit()
)

This comment has been minimized.

Copy link
@lohedges

lohedges Aug 30, 2023

Contributor

You're missing the except block here.


self._dyn = self._system.dynamics(
timestep=timestep,
Expand Down

0 comments on commit 6eea673

Please sign in to comment.