diff --git a/src/somd2/analysis/_analyse.py b/src/somd2/analysis/_analyse.py index 2d69a0b..fc5df02 100644 --- a/src/somd2/analysis/_analyse.py +++ b/src/somd2/analysis/_analyse.py @@ -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"): @@ -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") @@ -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 @@ -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() @@ -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 @@ -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() diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py index e031ab5..260f720 100644 --- a/src/somd2/runner/_runner.py +++ b/src/somd2/runner/_runner.py @@ -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: -------- @@ -208,7 +209,6 @@ def run_simulations(self): -------- results (list): List of simulation results. """ - print("here") import concurrent.futures as _futures results = [] @@ -307,10 +307,15 @@ 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": @@ -318,8 +323,8 @@ def _run(system, map, lambda_value, lam_minimisation=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: @@ -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: @@ -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, }, ) diff --git a/src/somd2/runner/_sire_merge_runsim.py b/src/somd2/runner/_sire_merge_runsim.py index f49b756..4ab317e 100644 --- a/src/somd2/runner/_sire_merge_runsim.py +++ b/src/somd2/runner/_sire_merge_runsim.py @@ -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() + ) self._dyn = self._system.dynamics( timestep=timestep,