diff --git a/README.md b/README.md index d5673a0..296cad1 100644 --- a/README.md +++ b/README.md @@ -25,30 +25,101 @@ Note: If you are using the system wide python environment instead of a virtual e Example script: ```python -import hercules as he +from hercules import KassLocustP3, SimConfig, ConfigList -sim = he.KassLocustP3('/path/to/your/workingDir') +sim = KassLocustP3('/path/to/your/workingDir', use_kass=True, use_locust=True) +configlist = ConfigList(info='Additional meta data') #just an example -config = he.SimConfig('yourSimulationName', phase='Phase3', kass_file_name='someXMLFile.xml', +config = SimConfig(phase='Phase3', kass_file_name='someXMLFile.xml', locust_file_name='someJSONFile.json', - nChannels=2, seedLocust=1, vRange=7.0, - eggFilename='someFileName.egg', seedKass =12534, xMin=-0.1e-5, - xMax=0.1e-5, tMax=0.5e-6, + n_channels=2, seed_locust=1, v_range=7.0, + egg_filename='someFileName.egg', seed_kass =12534, x_min=0.1e-5, + x_max=0.1e-5, t_max=0.5e-6, geometry='FreeSpaceGeometry_V00_00_10.xml') -sim(config) #can also take a list of configs +configlist.add_config(config) +#runs the simulations +sim(configlist) ``` -The example above runs a single phase 3 Kassiopeia-Locust simulation with the given parameters. All parameters except for the simulation name are optional. Omitted parameters in general take on default values with the exception being the seeds which are generated on the fly. The phase parameter can take the values 'Phase2' or 'Phase3' (default). Hercules generates the config files for Kassiopeia and Locust based on the inputs, the selected phase and the template config files which also provide the default values. All config files will be taken from the [hexbug](https://github.com/project8/hexbug/tree/459dffe30eea7d8bab9ddff78b63fda5198041ad) repository. Once hercules is installed you can run the script from anywhere specifying any working directory that you want and it will always be able to find the config files. Config files from hexbug (including Transfer functions and trap geometries) are passed by just their names as demonstrated above. Hercules will look for them in the appropriate directory of hexbug depending on the phase. In most cases you want to use the defaults for 'kass_file_name' and 'loucst_file_name'. +The example above runs a single phase 3 Kassiopeia-Locust simulation with the given parameters. Hercules is most useful if you run `config=...` and `configlist.add_config(config)` in an arbitrary python loop. The line `sim(configlist)` will always run all the simulation configurations from the list. All parameters are optional. Omitted parameters in general take on default values with the exception being the seeds which are generated on the fly. The phase parameter can take the values 'Phase2' or 'Phase3' (default). Hercules generates the config files for Kassiopeia and Locust based on the inputs, the selected phase and the template config files which also provide the default values. All config files will be taken from the [hexbug](https://github.com/project8/hexbug/tree/459dffe30eea7d8bab9ddff78b63fda5198041ad) repository. Once hercules is installed you can run the script from anywhere specifying any working directory that you want and it will always be able to find the config files. Config files from hexbug (including Transfer functions and trap geometries) are passed by just their names as demonstrated above. Hercules will look for them in the appropriate directory of hexbug depending on the phase. In most cases you want to use the defaults for 'kass_file_name' and 'loucst_file_name'. If you need the full list of simulation parameters you can ask hercules for a help message. `he.SimConfig.help()` will print a full list of all available keyword arguments with a short explanation for each one. -You can find example scripts in [examples](./examples). Hercules scripts work in a desktop environment as well as on the grace cluster without requiring any modifications. +You can find example scripts in [examples](./examples). Hercules scripts work in a desktop environment as well as on the grace cluster without requiring any modifications. + +Running hercules with the example from above will create a `hercules.Dataset` in the specified working directory. A `hercules.Dataset` is a dataformat in the form of an indexed directory structure which handles like a normal directory. For each configuration in the `configlist` hercules creates a subdirectory in the working directory that is called `run{i}` with `i` being the incremental number of configurations. In addition to that the directory contains a text file `info.txt` with some info about the dataset (meta data and parameter axes) and the very important file `index.he`. The latter is the pickled `hercules.Dataset` python object. Its core is a hashmap to map configuration parameters to the corresponding paths in the working directory. The class implements some utility for convenient recovery of any data stored in the subdirectories. In the example below a loop is used to recover all egg files but note that under each path other data than just the egg file can be found (with Locust as default you get at least log files and a json file with the configuration), which can be accessed with that path. See [dataset_example](./examples/dataset_example.py) for more details on how to utilize the `Dataset` class. + +```python +from hercules import Dataset, LocustP3File + +dataset = Dataset.load('/path/to/your/workingDir') + +for param, path in dataset: + data = LocustP3File(path / 'someFileName.egg') +``` + +Another interesting feature is the use of python scripts for post-processing. Any python script located in [hexbug/CRESana](./hercules/hexbug/CRESana/) (was implemented for the use with [CRESana](https://github.com/MCFlowMace/CRESana)) can be passed by its name and for each configuration it will be run after Kassiopeia and Locust. This represents another way of producing more data for a single configuration which can be retrieved via the path from the `Dataset`. + +```python +sim = KassLocustP3('/path/to/your/workingDir', use_kass=True, use_locust=True, python_script='post_processing.py') +``` + +The sole command line argument of these python scripts is the path of the result. More parameters can be used in the python script by importing the `SimConfig` object of the configuration via json. Thus the top of these scripts should be like + +```python +from hercules import SimConfig + +path = sys.argv[1] +config = SimConfig.from_json(path + '/SimConfig.json').to_dict() +#get pitch angle +pitch = config['kass-config']['theta_min'] + +``` + +Finally it is important to mention that the use of Kassiopeia and Locust is optional and can both be turned off as seen below. It depends on the configuration if it makes sense to run like that. + +```python +sim = KassLocustP3('/path/to/your/workingDir', use_kass=True, use_locust=False, python_script='run_no_locust.py') +``` + +Without both Locust and Kassiopeia hercules turns into a simple convenience tool for running python scripts on a parameter grid on the grace cluster with the `hercules.Dataset` as output. In this case the simpler and more flexible `SimpleSimConfig` can be used. It supports any configuration parameters as passed via keyword arguments, which can be passed to the script. Together with the metadata (`info`) which is added on creation of the `ConfigList` they will be represented in the final `Dataset`. Example: + +```python +sim = KassLocustP3('/path/to/your/workingDir', use_kass=False, use_locust=False, python_script='run_no_kass_no_locust.py') +configlist = ConfigList(info='Additional meta data') +config = SimpleSimConfig(x=2., some_exotic_data_name='interesting_value') +configlist.add_config(config) +sim(configlist) +``` + +Corresponding script head of `run_no_kass_no_locust.py`: + +```python +from hercules import SimpleSimConfig + +path = sys.argv[1] +config = SimpleSimConfig.from_json(path + '/SimConfig.json').to_dict() + +#get config parameters +x = config['x'] +some_exotic_data_name = config['some_exotic_data_name'] +``` + +## Running on grace cluster + +For running on the grace cluster there are a couple of extra keyword arguments for the `KassLocustP3` class. + +```python +sim(config_list, memory='1000', timelimit='01:00:00', n_cpus=8, batch_size=3) +``` + +Setting `timelimit` and `memory` only as high as required for the job is good practice and should theoretically help with job scheduling. `batch_size` determines how many entries in `config_list` are combined into a single job, you need to make sure to have a `batch_size` that gets you job run times >10 minutes since the cluster is not well suited for high job throughput. `n_cpus` defaults to 2 for the use with Locust and Kass+Locust alone does not profit from using more than that. Setting it to higher values is only useful if your postprocessing python script uses muliple processes. ## Tests -To test whether Docker is working on desktop, the `test_eggreader.py` provides a separate test. Run the following in cmd line: +To test whether Docker is working on your device, the unittest `test_locustP2.py` can be used. Run the following in cmd line: ```sh cd ./test -python -m unittest test_eggreader.EggReaderTest.test_locust +python -m unittest test_locustP2.py ``` diff --git a/examples/dataset_example.py b/examples/dataset_example.py new file mode 100644 index 0000000..70f0ef5 --- /dev/null +++ b/examples/dataset_example.py @@ -0,0 +1,50 @@ + +""" + +Author: F. Thomas +Date: August 07, 2023 + +""" + +from hercules import Dataset, LocustP3File +from pathlib import Path + +module_dir = Path(__file__).parent.absolute() +dataset_name = 'workingDirP3' + +path = module_dir / dataset_name + +#load the hercules dataset +dataset = Dataset.load(path) + +#print the metadata which is valid for all entries in the dataset +print(dataset.meta_data) + +#print the axes of the dataset grid +axes = dataset.axes +for i, ax in enumerate(axes): + print(f'{dataset.config_data_keys[i]} -> {ax}' ) + +#get path to data by index of all data axes +param, path = dataset.get_path([0, 0, 0, 0, 0], method='index') +print(param, path) + +#use path +data = LocustP3File(path / 'someFileName.egg') + +#get path to data by exact value +param, path = dataset.get_path([0., 0., 0., 87., 18600.], method='exact') +print(param, path) + +#get path to data by exact value +param, path = dataset.get_path([axes[0][3], axes[1][0], axes[2][0], axes[3][0], axes[4][0]], method='exact') +print(param, path) + +#get path to data by value using nearest neighbor interpolation +param, path = dataset.get_path([0.0, 0.0, 0.0, 0.0, 0.0], method='interpolated') +print(param, path) + +#iterate over all data in dataset +for param, path in dataset: + print(param, path) + data = LocustP3File(path / 'someFileName.egg') \ No newline at end of file diff --git a/examples/locustP2_example.py b/examples/locustP2_example.py index c920b24..eacf6ad 100644 --- a/examples/locustP2_example.py +++ b/examples/locustP2_example.py @@ -6,15 +6,15 @@ """ -import hercules as he +from hercules import KassLocustP3, SimConfig, ConfigList from pathlib import Path import numpy as np module_dir = Path(__file__).parent.absolute() #just an example -config_list = [] -sim = he.KassLocustP3(str(module_dir) + '/workingDir') +sim = KassLocustP3(str(module_dir) + '/workingDir', use_kass=True, use_locust=True) +configlist = ConfigList() n_r = 1 r_vals = np.linspace(0.0, 0.03, n_r) @@ -24,13 +24,12 @@ for i, r in enumerate(r_vals): for j, theta in enumerate(theta_vals): - config = he.SimConfig('someDirName_{0:1.3f}_{1:1.3f}'.format(r, theta), - phase='Phase2', egg_filename='someFileName.egg', + config = SimConfig( phase='Phase2', egg_filename='someFileName.egg', theta_min = theta, theta_max = theta, x_min = r, x_max = r, v_range=1.0e-6, record_size=20000, t_max=7.5e-5, geometry='Trap_3.xml') - config_list.append(config) + configlist.add_config(config) -sim(config_list) +sim(configlist) diff --git a/examples/locustP3_example.py b/examples/locustP3_example.py index a9c779a..881bcc7 100644 --- a/examples/locustP3_example.py +++ b/examples/locustP3_example.py @@ -6,24 +6,24 @@ """ -import hercules as he +from hercules import SimConfig, KassLocustP3, ConfigList from pathlib import Path import numpy as np module_dir = Path(__file__).parent.absolute() #just an example -config_list = [] -sim = he.KassLocustP3(str(module_dir) + '/workingDir') +sim = KassLocustP3(str(module_dir) + '/workingDirP3', use_kass=True, use_locust=True) +configlist = ConfigList() -simulations = 20 +simulations = 5 r_vals = np.linspace(0.0, 0.03, simulations) for i, r in enumerate(r_vals): - config = he.SimConfig('someDirName_{0:1.3f}'.format(r), n_channels=2, seed_locust=1, + config = SimConfig(n_channels=2, seed_locust=1, v_range=7.0, egg_filename='someFileName.egg', x_min=r, x_max=r, t_max=0.5e-6, geometry='FreeSpaceGeometry_V00_00_10.xml') - config_list.append(config) + configlist.add_config(config) -sim(config_list) +sim(configlist) diff --git a/hercules/__init__.py b/hercules/__init__.py index 542d164..4d517a3 100644 --- a/hercules/__init__.py +++ b/hercules/__init__.py @@ -7,6 +7,6 @@ """ from .simulation import KassLocustP3 -from .simconfig import SimConfig +from .simconfig import SimConfig, SimpleSimConfig, ConfigList from .eggreader import LocustP3File from .dataset import Dataset diff --git a/hercules/constants.py b/hercules/constants.py index 557105b..1b6b1e9 100644 --- a/hercules/constants.py +++ b/hercules/constants.py @@ -8,7 +8,7 @@ __all__ = [] -from pathlib import Path, PosixPath +from pathlib import Path, PurePosixPath from .configuration import Configuration @@ -16,8 +16,8 @@ HEXBUG_DIR = MODULE_DIR / 'hexbug' #container is running linux #-> make sure it's PosixPath when run from windows -HEXBUG_DIR_CONTAINER = PosixPath('/') / 'tmp' -OUTPUT_DIR_CONTAINER = PosixPath('/') / 'home' +HEXBUG_DIR_CONTAINER = PurePosixPath('/') / 'tmp' +OUTPUT_DIR_CONTAINER = PurePosixPath('/') / 'home' LOCUST_CONFIG_NAME_P2 = 'LocustPhase2Template.json' KASS_CONFIG_NAME_P2 = 'Project8Phase2_electrons.xml' LOCUST_CONFIG_NAME_P3 = 'LocustPhase3Template.json' diff --git a/hercules/dataset.py b/hercules/dataset.py index d213b3d..ee3d4f4 100644 --- a/hercules/dataset.py +++ b/hercules/dataset.py @@ -8,9 +8,8 @@ import numpy as np from scipy.interpolate import interp1d -import pickle +import dill as pickle from pathlib import Path -from math import sqrt, atan2 from .constants import PY_DATA_NAME @@ -24,63 +23,61 @@ def __init__(self, x): def __call__(self, x): return self.x + class Dataset: - __version = '1.0' + _class_version = '2.0' - def __init__(self, directory): + def __init__(self, directory, config_list): + + self._directory = Path(directory) + self._directory.mkdir(parents=True, exist_ok=True) + self._version = self._class_version + self._make_index(config_list) - self.directory = Path(directory) + def _make_index(self, config_list): + """Create the index dictionary. - def make_index(self, config_list): + Parameters + ---------- + config_list : ConfigList + A ConfigList object + """ print('Making file index') - self.index = {} - r_np = np.empty(len(config_list)) - phi_np = np.empty(len(config_list)) - z_np = np.empty(len(config_list)) - pitch_np = np.empty(len(config_list)) - energy_np = np.empty(len(config_list)) + self._index = {} + self._meta_data = config_list.get_meta_data() + self._config_data_keys = list(config_list.get_config_data_keys()) # maps a list index to a key + config_list_internal = config_list.get_internal_list() + + self._axes = [np.empty(len(config_list_internal)) for k in self._config_data_keys] - for i, sim_config in enumerate(config_list): + for i, sim_config in enumerate(config_list_internal): path = sim_config.sim_name - x = sim_config._kass_config._config_dict['x_min'] - y = sim_config._kass_config._config_dict['y_min'] - z = sim_config._kass_config._config_dict['z_min'] - pitch = sim_config._kass_config._config_dict['theta_min'] - energy = sim_config._kass_config._config_dict['energy'] - - r = sqrt(x**2 + y**2) - phi = atan2(y, x) - - self.index[energy, pitch, r, phi, z] = path - - r_np[i] = r - phi_np[i] = phi - z_np[i] = z - pitch_np[i] = pitch - energy_np[i] = energy + config_data = sim_config.get_config_data() + + for k in config_data: + k_ind = self._config_data_keys.index(k) + self._axes[k_ind][i] = config_data[k] - self.r = np.sort(np.unique(r_np)) - self.phi = np.sort(np.unique(phi_np)) - self.z = np.sort(np.unique(z_np)) - self.pitch = np.sort(np.unique(pitch_np)) - self.energy = np.sort(np.unique(energy_np)) + self._index[tuple(config_data.values())] = path + + for i in range(len(self._axes)): + self._axes[i] = np.sort(np.unique(self._axes[i])) - self.interpolate_all() + self._interpolate_all() - def interpolate_all(self): + def _interpolate_all(self): print('Making interpolation') + + self._axes_int = [] + + for ax in self._axes: + self._axes_int.append(self._interpolate(ax)) - self.r_int = self.interpolate(self.r) - self.phi_int = self.interpolate(self.phi) - self.z_int = self.interpolate(self.z) - self.pitch_int = self.interpolate(self.pitch) - self.energy_int = self.interpolate(self.energy) - - def interpolate(self, x): + def _interpolate(self, x): if len(x)>1: x_int = interp1d(x, x, kind='nearest', bounds_error=None, fill_value='extrapolate') @@ -89,43 +86,119 @@ def interpolate(self, x): return x_int - def get_data(self, energy, pitch, r, phi, z, interpolation=True): + def get_data(self, params, interpolation=True): - parameters, sim_path = self.get_path(energy, pitch, r, phi, z, interpolation=interpolation) + parameters, sim_path = self.get_path(params, interpolation=interpolation) - path = sim_path.relative_to(self.directory) + path = sim_path.relative_to(self._directory) return parameters, self.load_sim(path) - def load_sim(self, path): - return np.load(self.directory / path / PY_DATA_NAME) - - def get_path(self, energy, pitch, r, phi, z, interpolation=True): + def _load_sim(self, path): + return np.load(self._directory / path / PY_DATA_NAME) - if interpolation: - energy_i = self.energy_int(energy).item() - pitch_i = self.pitch_int(pitch).item() - r_i = self.r_int(r).item() - phi_i = self.phi_int(phi).item() - z_i = self.z_int(z).item() + def get_path(self, params, method='interpolated'): + + if len(params) != len(self._axes): + raise ValueError(f'params has len {len(params)} but dataset expects len {len(self._axes)}!') + + if method == 'interpolated': + key = [self._axes_int[i](params[i]).item() for i in range(len(params))] + elif method == 'index': + key = [self._axes[i][params[i]] for i in range(len(params))] + elif method == 'exact': + key = params else: - energy_i = energy - pitch_i = pitch - r_i = r - phi_i = phi - z_i = z - - parameters = (energy_i, pitch_i, r_i, phi_i, z_i) - sim_path = self.index[parameters] + raise ValueError("method can only take values 'interpolated', 'index' or 'exact'!") + + parameters = tuple(key) + + sim_path = self._index.get(parameters) # self._index[parameters] + + if sim_path is None: + raise KeyError(f'{parameters} is not part of the dataset!') - return parameters, self.directory / sim_path + return parameters, self._directory / sim_path + + def __iter__(self): + self._it_index = tuple(0 for i in range(len(self.shape))) + self._it_stop = False + return self + + def __next__(self): + + if not self._it_stop: + + new_index = list(self._it_index) + + for i in reversed(range(len(self._it_index))): + new_index[i] += 1 + if new_index[i] == self.shape[i]: + new_index[i] = 0 + else: + break + + old_index = self._it_index + self._it_index = tuple(new_index) + + if self._it_index == tuple(0 for i in range(len(self.shape))): + self._it_stop = True + + return self.get_path(old_index, method='index') + else: + raise StopIteration + + @property + def config_data_keys(self): + return self._config_data_keys + + @property + def axes(self): + return self._axes + + @property + def shape(self): + return tuple(len(ax) for ax in self._axes) + + @property + def meta_data(self): + return self._meta_data def dump(self): - pickle.dump(self, open(self.directory/'index.he', "wb"), protocol=4) + with open(self._directory/'index.he', "wb") as f: + pickle.dump(self, f, protocol=4) + + with open(self._directory/'info.txt', "w") as f: + f.write(f'Hercules dataset version {self._version}\n') + f.write('Metadata:\n') + f.write(str(self._meta_data)) + f.write('\n\n') + f.write('Dataset has following configurations:\n') + + for i, ax in enumerate(self._axes): + n = len(ax) + lower = ax[0] + upper = ax[-1] + ax_name = self._config_data_keys[i] + f.write(f'{ax_name}: {n} values in [{lower},{upper}] \n') @classmethod def load(cls, path): path_p = Path(path) - instance = pickle.load(open(path_p/'index.he', "rb")) - instance.directory = path_p - return instance + + with open(path_p/'index.he', "rb") as f: + instance = pickle.load(f) + + if type(instance) is not cls: + raise RuntimeError('Path does not point to a hercules dataset') + + if '_version' not in dir(instance): + instance_version = '1.0' + else: + instance_version = instance._version + + if instance_version != cls._class_version: + raise RuntimeError(f'Tried to load a version {instance_version} hercules dataset with version {cls._class_version}! To open this file you need an older hercules release') + + instance._directory = path_p + return instance \ No newline at end of file diff --git a/hercules/hexbug b/hercules/hexbug index 324aa54..edc4ed8 160000 --- a/hercules/hexbug +++ b/hercules/hexbug @@ -1 +1 @@ -Subproject commit 324aa540d65df531fb1ea45056bb6810f488e411 +Subproject commit edc4ed8d05d43d03b9c0552636af90901d3d9b21 diff --git a/hercules/simconfig.py b/hercules/simconfig.py index a03fd04..3365001 100644 --- a/hercules/simconfig.py +++ b/hercules/simconfig.py @@ -11,9 +11,8 @@ import time import json import re -from pathlib import Path, PosixPath -from abc import ABC, abstractmethod from copy import deepcopy +from math import sqrt, atan2 from .constants import (HEXBUG_DIR, HEXBUG_DIR_CONTAINER, OUTPUT_DIR_CONTAINER, LOCUST_CONFIG_NAME_P2, KASS_CONFIG_NAME_P2, @@ -983,7 +982,7 @@ class SimConfig: Name of the simulation """ - def __init__(self, sim_name, phase = 'Phase3', kass_file_name = None, + def __init__(self, phase = 'Phase3', kass_file_name = None, kass_unknown_args_translation = {}, locust_file_name = None, locust_unknown_args_translation = {}, **kwargs): @@ -1036,8 +1035,9 @@ def __init__(self, sim_name, phase = 'Phase3', kass_file_name = None, If phase is not 'Phase2' or 'Phase3'. """ - self._sim_name = sim_name + self._sim_name = None self._phase = phase + self._extra_meta_data = {} self._locust_config = LocustConfig(phase = phase, locust_file_name = locust_file_name, @@ -1084,6 +1084,10 @@ def _trigger_unknown_parameter_warnings(self, kwargs): def sim_name(self): return self._sim_name + @sim_name.setter + def sim_name(self, sim_name): + self._sim_name = sim_name + def to_json(self, file_name): """Write a json file with the entire simulation configuration.""" @@ -1126,10 +1130,10 @@ def from_json(cls, file_name): with open(file_name, 'r') as infile: config = json.load(infile) - sim_name = config['sim-name'] phase = config['phase'] - instance = cls(sim_name, phase=phase) + instance = cls(phase=phase) + instance.sim_name = config['sim-name'] instance._locust_config._config_dict = config['locust-config'] instance._kass_config._config_dict = config['kass-config'] @@ -1154,8 +1158,18 @@ def help(cls): print() print('Note that all keyword arguments are optional and take default values from the config files!') - def make_config_file(self, filename_locust, filename_kass): - """Create the final Kassiopeia and Locust config files. + def make_kass_config_file(self, filename_kass): + """Create the final Kassiopeia file. + + Parameters + ---------- + filename_kass : str + the path to the output Kassiopeia config file + """ + self._kass_config.make_config_file(filename_kass) + + def make_locust_config_file(self, filename_locust, filename_kass): + """Create the final Locust config file. Parameters ---------- @@ -1164,7 +1178,247 @@ def make_config_file(self, filename_locust, filename_kass): filename_kass : str the path to the output Kassiopeia config file """ - self._locust_config.set_xml(filename_kass) self._locust_config.make_config_file(filename_locust) - self._kass_config.make_config_file(filename_kass) + + def get_meta_data(self): + + #maybe incomplete + #add more when you realize you need more metadata + + tf_file_name = self._locust_config._config_dict[LocustConfig._array_signal_key].get(LocustConfig._tf_receiver_filename_key) + n_channels = self._locust_config._config_dict[LocustConfig._sim_key].get(LocustConfig._n_channels_key) + acq_rate = self._locust_config._config_dict[LocustConfig._sim_key].get(LocustConfig._acq_rate_key) + lo_f = self._locust_config._config_dict[LocustConfig._array_signal_key].get(LocustConfig._lo_frequency_key) + + meta_data = {} + meta_data.update(self._extra_meta_data) + + meta_data.update({'trap': self._kass_config._config_dict['geometry']}) + + if tf_file_name is not None: + meta_data.update({'transfer-function': tf_file_name}) + + if n_channels is not None: + meta_data.update({'n-channels': n_channels}) + + if acq_rate is not None: + meta_data.update({'acquisition-rate': acq_rate}) + + if lo_f is not None: + meta_data.update({'lo-frequency': lo_f}) + + return meta_data + + def add_meta_data(self, meta_data): + self._extra_meta_data = meta_data + + def get_config_data(self): + + config_data = {} + + x_min = self._kass_config._config_dict['x_min'] + y_min = self._kass_config._config_dict['y_min'] + z_min = self._kass_config._config_dict['z_min'] + pitch_min = self._kass_config._config_dict['theta_min'] + + x_max = self._kass_config._config_dict['x_max'] + y_max = self._kass_config._config_dict['y_max'] + z_max = self._kass_config._config_dict['z_max'] + pitch_max = self._kass_config._config_dict['theta_max'] + + energy = self._kass_config._config_dict['energy'] + + r_min = sqrt(x_min**2 + y_min**2) + phi_min = atan2(y_min, x_min) + + r_max = sqrt(x_max**2 + y_max**2) + phi_max = atan2(y_max, x_max) + + if x_min==x_max and y_min==y_max and z_min==z_max and pitch_min==pitch_max: + config_data['r'] = r_min + config_data['phi'] = phi_min + config_data['z'] = z_min + config_data['pitch'] = pitch_min + else: + config_data['r_min'] = r_min + config_data['phi_min'] = phi_min + config_data['z_min'] = z_min + config_data['pitch_min'] = pitch_min + config_data['r_max'] = r_max + config_data['phi_max'] = phi_max + config_data['z_max'] = z_max + config_data['pitch_max'] = pitch_max + + config_data['energy'] = energy + + return config_data + + +class SimpleSimConfig: + """A class for a more general simulation configuration + + This class is intended for the use with more general python scripts. + It supports an arbitrary number of keyword arguments. Arguments with the prefix 'meta_' + are treated as meta parameters. Meta parameters are expected to be the same for an entire dataset. + All other parameters are considered regular configuration parameters that should vary over a dataset. + + Attributes + ---------- + sim_name : str + Name of the simulation + """ + + def __init__(self, **kwargs): + """ + Parameters + ---------- + sim_name : str + Name of the simulation + **kwargs : + Arbitrary number of keyword arguments. + + """ + + self._sim_name = None + self._extract_kwargs(kwargs) + + def _extract_kwargs(self, kwargs): + + self._meta_data = {} + self._config_data = {} + #for e in kwargs: + # if e.startswith('meta_'): + # new_key = e.removeprefix('meta_') + # self._meta_data[new_key] = kwargs[e] + # else: + # self._config_data[e] = kwargs[e] + self._config_data = kwargs + + @property + def sim_name(self): + return self._sim_name + + @sim_name.setter + def sim_name(self, sim_name): + self._sim_name = sim_name + + def to_json(self, file_name): + """Write a json file with the entire simulation configuration.""" + + with open(file_name, 'w') as outfile: + json.dump({ 'sim-name': self._sim_name, + 'meta-data': self._meta_data, + 'config-data': self._config_data}, outfile, + indent=2)#, default=lambda x: x.config_dict) + + + def to_dict(self): + """Return a dictionary with the entire simulation configuration. + + Returns + ------- + dict + Nested dictionary with the simulation configuration + """ + + return {'sim-name': self._sim_name, + 'meta-data': self._meta_data, + 'config-data': self._config_data} + + @classmethod + def from_json(cls, file_name): + """Return a SimpleSimConfig from a json file. + + Creates a new instance of a SimpleSimConfig from the contents of a json file. + This should only be used with a json file that was created by the + `to_json` method. No checks applied for the validity of the json file. + + Returns + ------- + SimpleSimConfig + The new SimpleSimConfig instance + """ + + with open(file_name, 'r') as infile: + config = json.load(infile) + + instance = cls() + instance.sim_name = config['sim-name'] + instance._meta_data = config['meta-data'] + instance._config_data = config['config-data'] + + return instance + + @classmethod + def help(cls): + """Print documentation about the SimConfig. + + Prints the docstrings of the class and the __init__ method as well as + additional information about the accepted parameters in the keyword + arguments. The latter is provided via the two wrapped configurations. + + """ + print(cls.__doc__) + print(cls.__init__.__doc__) + + def get_meta_data(self): + return self._meta_data + + def get_config_data(self): + return self._config_data + + def add_meta_data(self, meta_data): + self._meta_data = meta_data + + +class ConfigList: + + def __init__(self, **kwargs): + self._config_list = [] + self._meta_data = kwargs + self._extra_meta_data = None + self._config_list_type = None + self._config_data_keys = None + + def add_config(self, config): + + n = len(self._config_list) + + if n == 0: + + common_keys = set(self._meta_data.keys()).intersection(config.get_meta_data().keys()) + + if len(common_keys)>0: + print('Warning, adding a config with metadata that overwrites an existing metadata entry! This might not be what you want.') + + self._extra_meta_data = config.get_meta_data() + self._meta_data.update(config.get_meta_data()) + self._config_list_type = type(config) + self._config_data_keys = config.get_config_data().keys() + + if type(config) is not self._config_list_type: + raise TypeError('All configurations in the configuration list have to be of the same type!') + + if config.get_meta_data() != self._extra_meta_data: + raise RuntimeError('All configurations in the configuration list need the same metadata') + + if config.get_config_data().keys() != self._config_data_keys: + raise RuntimeError('All configurations in the configuration list need the same configuration data keys') + + config.add_meta_data(self._meta_data) + config.sim_name = f'run{n}' + self._config_list.append(config) + + def get_internal_list(self): + return self._config_list + + def get_list_type(self): + return self._config_list_type + + def get_meta_data(self): + return self._meta_data + + def get_config_data_keys(self): + return self._config_data_keys + diff --git a/hercules/simulation.py b/hercules/simulation.py index ca21352..d8865c1 100644 --- a/hercules/simulation.py +++ b/hercules/simulation.py @@ -8,7 +8,7 @@ __all__ = ['KassLocustP3'] -from pathlib import Path, PosixPath +from pathlib import Path, PurePosixPath import subprocess from abc import ABC, abstractmethod import concurrent.futures as cf @@ -16,7 +16,7 @@ from math import sqrt, atan2 import pickle -from hercules.simconfig import SimConfig +from hercules.simconfig import ConfigList, SimpleSimConfig from .dataset import Dataset from .constants import (HEXBUG_DIR, HEXBUG_DIR_CONTAINER, OUTPUT_DIR_CONTAINER, LOCUST_CONFIG_NAME, KASS_CONFIG_NAME, SIM_CONFIG_NAME, @@ -160,8 +160,8 @@ class AbstractKassLocustP3(ABC): """An abstract base class for all KassLocust simulations.""" #configuration parameters - _p8_locust_dir = PosixPath(CONFIG.locust_path) / CONFIG.locust_version - _p8_compute_dir = PosixPath(CONFIG.p8compute_path) / CONFIG.p8compute_version + _p8_locust_dir = PurePosixPath(CONFIG.locust_path) / CONFIG.locust_version + _p8_compute_dir = PurePosixPath(CONFIG.p8compute_path) / CONFIG.p8compute_version def __init__(self, working_dir, use_locust=True, use_kass=False, python_script=None, direct=True): @@ -184,17 +184,16 @@ def __call__(self, config_list, **kwargs): Parameters ---------- - config_list : list - A list of SimConfig objects + config_list : ConfigList + A ConfigList object """ self.make_index(config_list) - self.run(config_list, **kwargs) + self.run(config_list.get_internal_list(), **kwargs) def make_index(self, config_list): - dataset = Dataset(self._working_dir) - dataset.make_index(config_list) + dataset = Dataset(self._working_dir, config_list) dataset.dump() @abstractmethod @@ -248,7 +247,7 @@ def factory(name, working_dir, use_locust=True, use_kass=False, class KassLocustP3Desktop(AbstractKassLocustP3): """A class for running KassLocust on a desktop.""" - _working_dir_container = PosixPath('/') / 'workingdir' + _working_dir_container = PurePosixPath('/') / 'workingdir' _command_script_name = 'locustcommands.sh' _container = CONFIG.container _max_workers = int(CONFIG.desktop_parallel_jobs) @@ -286,7 +285,7 @@ def run(self, sim_config_list, **kwargs): for future in tqdm(cf.as_completed(futures), total=len(futures)): future.result() - def _submit(self, sim_config: SimConfig): + def _submit(self, sim_config): #Submit the job with the given SimConfig #Creates all the necessary configuration files, directories and the #json output @@ -298,8 +297,13 @@ def _submit(self, sim_config: SimConfig): kass_file = output_dir / KASS_CONFIG_NAME config_dump = output_dir / SIM_CONFIG_NAME - sim_config.make_config_file(locust_file, kass_file) sim_config.to_json(config_dump) + + if self._use_locust: + sim_config.make_locust_config_file(locust_file, kass_file) + + if self._use_kass: + sim_config.make_kass_config_file(kass_file) if self._use_locust or self._use_kass: self._gen_command_script(output_dir) @@ -430,6 +434,7 @@ def _submit_job(self, **kwargs): module = 'module load dSQ;' n_cpus = 2 if self._use_locust else 1 + n_cpus = n_cpus if 'n_cpus' not in kwargs else kwargs['n_cpus'] memory = CONFIG.job_memory if 'memory' not in kwargs else kwargs['memory'] timelimit = CONFIG.job_timelimit if 'timelimit' not in kwargs else kwargs['timelimit'] @@ -459,8 +464,16 @@ def _add_job(self, sim_config): kass_file = output_dir / KASS_CONFIG_NAME config_dump = output_dir / SIM_CONFIG_NAME - sim_config.make_config_file(locust_file, kass_file) sim_config.to_json(config_dump) + + if self._use_locust: + sim_config.make_locust_config_file(locust_file, kass_file) + + if self._use_kass: + sim_config.make_kass_config_file(kass_file) + + if self._use_locust or self._use_kass: + self._gen_locust_script(output_dir) self._gen_locust_script(output_dir) cmd = self._assemble_command(output_dir) @@ -559,6 +572,9 @@ def __init__(self, working_dir, use_locust=True, use_kass=False, working_dir : str The string for the path to the working directory """ + + self._use_locust = use_locust + self._use_kass = use_kass self._kass_locust = AbstractKassLocustP3.factory(CONFIG.env, working_dir, @@ -574,7 +590,11 @@ def __call__(self, config_list, **kwargs): config_list : list or SimConfig Either a single SimConfig object or a list """ - if type(config_list) is not list: - config_list = [config_list] + if type(config_list) is not ConfigList: + raise TypeError('Needs an instance of ConfigList') + + if config_list.get_list_type() is SimpleSimConfig and (self._use_kass or self._use_locust): + raise TypeError('SimpleSimConfig is not compatible with the use of Locust or Kassiopeia!') + return self._kass_locust(config_list, **kwargs) diff --git a/setup.py b/setup.py index 4f85a08..f64f2e4 100644 --- a/setup.py +++ b/setup.py @@ -14,12 +14,13 @@ "h5py", "numpy", "scipy", - "tqdm" + "tqdm", + "dill" ] setuptools.setup( name="hercules", - version="0.4.3", + version="0.5.0", author="Florian Thomas, Mingyu (Charles) Li", author_email="fthomas@uni-mainz.de, mingyuli@mit.edu", description="https://github.com/project8/hercules", diff --git a/test/test_config.py b/test/test_config.py index 47e9853..dfa1b3a 100644 --- a/test/test_config.py +++ b/test/test_config.py @@ -6,33 +6,158 @@ """ from pathlib import Path -FILE_DIR = Path(__file__).parent.absolute() + +module_dir = Path(__file__).parent.absolute() import unittest -import hercules as he +from hercules import SimConfig, SimpleSimConfig, ConfigList import numpy as np -class ConfigTest(unittest.TestCase): - def test_generate_configs(self) -> None: - # Generates a list of config +class SimConfigTest(unittest.TestCase): + + def setUp(self) -> None: + + n_channels = 3 + x = 1. + y = 0. + theta = 90. + + self.config = SimConfig( + n_channels=n_channels, + seed_locust=42, + seed_kass=43, + egg_filename="simulation.egg", + x_min=x, + x_max=x, + y_min=y, + y_max=y, + z_min=0, + z_max=0, + theta_min=theta, + theta_max=theta, + t_max=5e-6, + v_range=3.0e-7, + geometry='FreeSpaceGeometry_V00_00_10.xml') + + self.file_name_locust = module_dir / 'locust.json' + self.file_name_kass = module_dir / 'kass.xml' + self.file_name_json = module_dir / 'test.json' + + def tearDown(self) -> None: + + if self.file_name_json.exists(): + self.file_name_json.unlink() + + if self.file_name_locust.exists(): + self.file_name_locust.unlink() + + if self.file_name_kass.exists(): + self.file_name_kass.unlink() + + def test_meta_data(self): + expected = {'trap': '[config_path]/Trap/FreeSpaceGeometry_V00_00_10.xml', + 'transfer-function': '/tmp/Phase3/TransferFunctions/FiveSlotTF.txt', + 'n-channels': 3, + 'lo-frequency': 25878100000.0} + + self.assertTrue(expected==self.config.get_meta_data()) + + def test_config_data(self): + expected = {'r': 1.0, 'phi': 0.0, 'z': 0, 'pitch': 90.0, 'energy': 18600.0} + + self.assertTrue(expected==self.config.get_config_data()) + + def test_to_json(self): + + self.config.to_json(self.file_name_json) + config_loaded = SimConfig.from_json(self.file_name_json) + self.assertTrue(config_loaded.to_dict()==self.config.to_dict()) + + def test_make_config_file(self): + + self.config.make_kass_config_file(self.file_name_kass) + self.config.make_locust_config_file(self.file_name_locust, self.file_name_kass) + + self.assertTrue(self.file_name_kass.exists()) + self.assertTrue(self.file_name_locust.exists()) + + def test_add_metadata(self): + + additional_meta_data = {'trap': 'this should be overwritten', + 'info': 'this is some additional info', + 'acquisition-rate': 'this will not be overwritten because SimConfig does not set it. Adding metadata CANNOT overwrite the config!'} + + expected = {'trap': '[config_path]/Trap/FreeSpaceGeometry_V00_00_10.xml', + 'info': 'this is some additional info', + 'acquisition-rate': 'this will not be overwritten because SimConfig does not set it. Adding metadata CANNOT overwrite the config!', + 'transfer-function': '/tmp/Phase3/TransferFunctions/FiveSlotTF.txt', + 'n-channels': 3, + 'lo-frequency': 25878100000.0} + + self.config.add_meta_data(additional_meta_data) + self.assertTrue(self.config.get_meta_data()==expected) + +class SimpleSimConfigTest(unittest.TestCase): + + def setUp(self) -> None: + + self.config = SimpleSimConfig(x=1., y=2., z=3.) + self.file_name_json = module_dir / 'test.json' + + def tearDown(self) -> None: + + if self.file_name_json.exists(): + self.file_name_json.unlink() + + def test_meta_data(self): + expected = {} + + self.assertTrue(expected==self.config.get_meta_data()) + + def test_config_data(self): + expected = {'x': 1., 'y': 2., 'z': 3.} + + self.assertTrue(expected==self.config.get_config_data()) + + def test_to_json(self): + + self.config.to_json(self.file_name_json) + config_loaded = SimpleSimConfig.from_json(self.file_name_json) + self.assertTrue(config_loaded.to_dict()==self.config.to_dict()) + + def test_add_metadata(self): + + meta_data = {'info1': 2, + 'info2': 'this is some additional info'} + + self.config.add_meta_data(meta_data) + self.assertTrue(self.config.get_meta_data()==meta_data) + + +class ConfigListTest(unittest.TestCase): + + def setUp(self) -> None: + n_channels = 3 + self.configlist = ConfigList(acquisition_rate=1., info='hello', trap='nonsense') + self.configlist_simple = ConfigList(acquisition_rate=1., info='hello simple', trap='nonsense', n_channels=n_channels) + + self.configlist_l = [] + self.configlist_simple_l = [] + r_range = np.linspace(0.002, 0.008, 8) theta_range = np.linspace(89.7, 90.0, 30) r_phi_range = np.linspace(0, 2 * np.pi / n_channels, 1) - config_list = [] - for theta in theta_range: for r_phi in r_phi_range: for r in r_range: x = r * np.cos(r_phi) y = r * np.sin(r_phi) - r_phi_deg = np.rad2deg(r_phi) - name = "Sim_theta_{:.4f}_R_{:.4f}_phi_{:.4f}".format( - theta, r, r_phi_deg) - config = he.SimConfig(name, + + sim_config = SimConfig( n_channels=n_channels, seed_locust=42, seed_kass=43, @@ -48,31 +173,104 @@ def test_generate_configs(self) -> None: t_max=5e-6, v_range=3.0e-7, geometry='FreeSpaceGeometry_V00_00_10.xml') + + self.configlist_l.append(sim_config) - json_file = FILE_DIR/'SimConfig.json' - locust_file = FILE_DIR/'Locust.json' - # Note kass file name has to be the following for dict comparison - kass_file = FILE_DIR/'LocustKassElectrons.xml' + simple_config = SimpleSimConfig(r=r, phi=r_phi, z=0., energy=0., pitch=90.) + self.configlist_simple_l.append(simple_config) - config.to_json(json_file) + def test_add_config(self): + + self.configlist.add_config(self.configlist_l[0]) + self.configlist.add_config(self.configlist_l[1]) - config2 = he.SimConfig.from_json(json_file) - config2.make_config_file(locust_file, kass_file) - self.assertDictEqual(config.to_dict(), config2.to_dict()) + #adding wrong type + with self.assertRaises(TypeError) as cm: + self.configlist.add_config(self.configlist_simple_l[0]) - config_list.append(config) + self.configlist_simple.add_config(self.configlist_simple_l[0]) + self.configlist_simple.add_config(self.configlist_simple_l[1]) - def tearDown(self) -> None: - # Clean up - files_json = FILE_DIR.glob('*.json') - files_xml = FILE_DIR.glob('*.xml') - filtered_files = sorted(files_json) + sorted(files_xml) + #adding non matching config data + with self.assertRaises(RuntimeError) as cm: + self.configlist_simple.add_config(SimpleSimConfig(x='foo')) + + def test_type(self): + + self.configlist.add_config(self.configlist_l[0]) + + self.configlist_simple.add_config(self.configlist_simple_l[0]) + + self.assertEqual(self.configlist.get_list_type(), SimConfig) + self.assertEqual(self.configlist_simple.get_list_type(), SimpleSimConfig) + + + def add_all(self): + for config in self.configlist_l: + self.configlist.add_config(config) + + for config in self.configlist_simple_l: + self.configlist_simple.add_config(config) + + def test_add_all(self): + + self.add_all() + + self.assertEqual(len(self.configlist.get_internal_list()), len(self.configlist_l)) + self.assertEqual(len(self.configlist_simple.get_internal_list()), len(self.configlist_simple_l)) + + def test_meta_data(self): + + self.add_all() + + expected = {'acquisition_rate': 1.0, + 'info': 'hello', + 'trap': '[config_path]/Trap/FreeSpaceGeometry_V00_00_10.xml', + 'transfer-function': '/tmp/Phase3/TransferFunctions/FiveSlotTF.txt', + 'n-channels': 3, + 'lo-frequency': 25878100000.0} + + self.assertEqual(expected, self.configlist.get_meta_data()) + + for config in self.configlist.get_internal_list(): + self.assertEqual(expected, config.get_meta_data()) + + expected = {'acquisition_rate': 1.0, + 'info': 'hello simple', + 'trap': 'nonsense', + 'n_channels': 3} + + self.assertEqual(expected, self.configlist_simple.get_meta_data()) + + for config in self.configlist_simple.get_internal_list(): + self.assertEqual(expected, config.get_meta_data()) + + def test_config_data(self): + + self.add_all() + + expected = ['energy', 'phi', 'pitch', 'r', 'z'] + + self.assertEqual(sorted(list(self.configlist.get_config_data_keys())), expected) + self.assertEqual(sorted(list(self.configlist_simple.get_config_data_keys())), expected) + + for config in self.configlist.get_internal_list(): + config_data_keys = config.get_config_data().keys() + self.assertEqual(sorted(list(config_data_keys)), expected) + + for config in self.configlist_simple.get_internal_list(): + config_data_keys = config.get_config_data().keys() + self.assertEqual(sorted(list(config_data_keys)), expected) + + def test_names(self): + + self.add_all() - for f in filtered_files: - # Note f is a Path object - # Deletes the file - f.unlink() + for i, config in enumerate(self.configlist.get_internal_list()): + self.assertEqual(config.sim_name,f'run{i}') + for i, config in enumerate(self.configlist_simple.get_internal_list()): + self.assertEqual(config.sim_name,f'run{i}') if __name__ == '__main__': unittest.main() diff --git a/test/test_dataset.py b/test/test_dataset.py new file mode 100644 index 0000000..a7813f6 --- /dev/null +++ b/test/test_dataset.py @@ -0,0 +1,112 @@ + +""" + +Author: F. Thomas +Date: Aug 05, 2023 + +""" + +from hercules import SimpleSimConfig, Dataset, ConfigList +from pathlib import Path +import unittest +import numpy as np +import shutil + +module_dir = Path(__file__).parent.absolute() +test_dataset_name = 'test_directory' +test_path = module_dir / test_dataset_name + +class DatasetTest(unittest.TestCase): + + def setUp(self) -> None: + + clist = ConfigList(sr=200., info='hello') + + y = 3. + for x in range(10): + for z in range(5,7,1): + clist.add_config(SimpleSimConfig(x=x, y=y, z=z)) + + self.d = Dataset(test_path, clist) + + def tearDown(self) -> None: + shutil.rmtree(test_path) + + def test_axes(self) -> None: + expected_result = [np.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]), np.array([3.]), np.array([5., 6.])] + axes = self.d.axes + equals = min([np.array_equal(axes[i], expected_result[i]) for i in range(len(expected_result))]) + self.assertTrue(equals) + + def test_config_data_keys(self) -> None: + expected_result = ['x', 'y', 'z'] + self.assertTrue(expected_result==self.d.config_data_keys) + + def test_shape(self) -> None: + expected_result = (10, 1, 2) + self.assertTrue(expected_result==self.d.shape) + + def test_meta_data(self) -> None: + expected_result = {'sr': 200.0, 'info': 'hello'} + self.assertTrue(expected_result==self.d.meta_data) + + def test_get_path(self) -> None: + + expected_result_1 = ((9.0, 3.0, 6.0), (test_path / 'run19').absolute()) + expected_result_2 = ((1.0, 3.0, 5.0), (test_path /'run2').absolute()) + expected_result_3 = ((4.0, 3.0, 6.0), (test_path / 'run9').absolute()) + + self.assertTrue(expected_result_1==self.d.get_path([100, 100, 100], method='interpolated')) + self.assertTrue(expected_result_2==self.d.get_path([1., 3., 5.], method='exact')) + self.assertTrue(expected_result_3==self.d.get_path([4, 0, 1], method='index')) + + with self.assertRaises(ValueError) as cm: + self.d.get_path([4, 0, 1], method='inde') + + def test_iterator(self) -> None: + + expected_result = [((0.0, 3.0, 5.0), (test_path / 'run0').absolute()), + ((0.0, 3.0, 6.0), (test_path / 'run1').absolute()), + ((1.0, 3.0, 5.0), (test_path / 'run2').absolute()), + ((1.0, 3.0, 6.0), (test_path / 'run3').absolute()), + ((2.0, 3.0, 5.0), (test_path / 'run4').absolute()), + ((2.0, 3.0, 6.0), (test_path / 'run5').absolute()), + ((3.0, 3.0, 5.0), (test_path / 'run6').absolute()), + ((3.0, 3.0, 6.0), (test_path / 'run7').absolute()), + ((4.0, 3.0, 5.0), (test_path / 'run8').absolute()), + ((4.0, 3.0, 6.0), (test_path / 'run9').absolute()), + ((5.0, 3.0, 5.0), (test_path / 'run10').absolute()), + ((5.0, 3.0, 6.0), (test_path / 'run11').absolute()), + ((6.0, 3.0, 5.0), (test_path / 'run12').absolute()), + ((6.0, 3.0, 6.0), (test_path / 'run13').absolute()), + ((7.0, 3.0, 5.0), (test_path / 'run14').absolute()), + ((7.0, 3.0, 6.0), (test_path / 'run15').absolute()), + ((8.0, 3.0, 5.0), (test_path / 'run16').absolute()), + ((8.0, 3.0, 6.0), (test_path / 'run17').absolute()), + ((9.0, 3.0, 5.0), (test_path / 'run18').absolute()), + ((9.0, 3.0, 6.0), (test_path / 'run19').absolute())] + + result = [] + + for entry in self.d: + result.append(entry) + + self.assertTrue(result==expected_result) + + def test_dump_load(self) -> None: + + self.d.dump() + d = Dataset.load(test_path) + + axes_self = self.d.axes + axes_load = d.axes + equals = min([np.array_equal(axes_self[i], axes_load[i]) for i in range(len(axes_self))]) + self.assertTrue(equals) + self.assertTrue(self.d.meta_data == d.meta_data) + self.assertTrue(self.d.config_data_keys == d.config_data_keys) + self.assertTrue(self.d._directory == d._directory) + self.assertTrue(self.d._index == d._index) + + +if __name__ == '__main__': + unittest.main() diff --git a/test/test_eggreader.py b/test/test_eggreader.py index 65cd522..ea4a2c9 100644 --- a/test/test_eggreader.py +++ b/test/test_eggreader.py @@ -1,6 +1,6 @@ """ -Author: Mingyu (Charles) Li +Author: Mingyu (Charles) Li, Florian Thomas Date: Apr. 12, 2021 """ @@ -10,39 +10,43 @@ from matplotlib.figure import Figure -FILE_DIR = Path(__file__).parent.absolute() - import matplotlib import unittest -import hercules as he +from hercules import KassLocustP3, SimConfig, ConfigList, Dataset, LocustP3File import numpy as np import matplotlib.pyplot as plt +import shutil + +#matplotlib.use("Agg") # No GUI -matplotlib.use("Agg") # No GUI +module_dir = Path(__file__).parent.absolute() +test_dataset_name = 'egg_reader_test' +test_path = module_dir / test_dataset_name +egg_filename = 'simulation.egg' class EggReaderTest(unittest.TestCase): - def setUp(self) -> None: + + @classmethod + def setUpClass(cls) -> None: n_ch = 3 r_range = np.linspace(0.000, 0.010, 1) - theta_range = np.linspace(85.0, 90.0, 1) + theta_range = np.linspace(85.0, 90.0, 2) r_phi_range = np.linspace(0, 2 * np.pi / n_ch, 1) - test_data_dict = {} + sim = KassLocustP3(test_path, use_kass=True, use_locust=True) + configlist = ConfigList() + for theta in theta_range: for r_phi in r_phi_range: for r in r_range: x = r * np.cos(r_phi) y = r * np.sin(r_phi) - r_phi_deg = np.rad2deg(r_phi) - name = "Sim_theta_{:.4f}_R_{:.4f}_phi_{:.4f}".format( - theta, r, r_phi_deg) - config = he.SimConfig( - name, + config = SimConfig( n_channels=n_ch, seed_locust=42, seed_kass=43, - egg_filename="simulation.egg", + egg_filename=egg_filename, x_min=x, x_max=x, y_min=y, @@ -57,44 +61,40 @@ def setUp(self) -> None: lo_frequency=25.8881e9, acq_rate=250.0, geometry='FreeSpaceGeometry_V00_00_10.xml') - test_data_dict[name] = config - - self.test_data_dict = test_data_dict - - # Check if the test dir exists and names are correct - self.test_data_dir = test_data_dir = FILE_DIR / "test_dir" - test_data_dir.mkdir(exist_ok=True) - sub_dir_list = [ - d.parts[-1] for d in test_data_dir.iterdir() if d.is_dir() - ] - missing_dir_list = list( - set(self.test_data_dict.keys()) - set(sub_dir_list)) - print( - "The following test data are missing: {}".format(missing_dir_list)) - - # Create missing test data - # Note: if data is missing, must run the tests in cmd line to create the data - if len(missing_dir_list) != 0: - print("Creating missing data...") - missing_data_config = [test_data_dict[l] for l in missing_dir_list] - sim = he.KassLocustP3(str(test_data_dir)) - sim(missing_data_config) - - def test_locust(self) -> None: - # Test sim generation (done in setUp), substitutes regular locust test - # This test is always true. - self.assertTrue(True) + configlist.add_config(config) + + sim(configlist) + + cls.dataset = Dataset.load(test_path) + + @classmethod + def tearDownClass(cls) -> None: + shutil.rmtree(test_path) + + def test_0(self) -> None: + + paths =[test_path / 'index.he', + test_path / 'info.txt', + test_path / 'run0' / 'simulation.egg', + test_path / 'run1' / 'simulation.egg' + ] + print('Testing simulation successful') + for path in paths: + self.assertTrue(path.exists()) def test_ts_stream(self) -> None: - for k in self.test_data_dict.keys(): - self._load_ts_stream(k) - self._quick_load_ts_stream(k) + for _, path in self.dataset: + print(path) + self._load_ts_stream(path) + self._quick_load_ts_stream(path) + + ok = input('Plots look ok? (y/n): ').lower().strip() == 'y' + self.assertTrue(ok) def _load_ts_stream(self, name) -> None: # Plot some specific test data - file_name = self.test_data_dir.joinpath(name).joinpath( - "simulation.egg") - file = he.LocustP3File(str(file_name)) + + file = LocustP3File(str(name / egg_filename)) n_streams = file.n_streams for s in range(n_streams): @@ -122,13 +122,12 @@ def _load_ts_stream(self, name) -> None: ax.set_xlabel(r"Time $[s]$") ax.set_ylabel(r"DAQ V $[V]$") ax.legend(loc="best") - self._save_fig(fig, title) + #self._save_fig(fig, f'{name}/plot_{s}.pdf', title) + plt.show() def _quick_load_ts_stream(self, name) -> None: # Plot some specific test data - file_name = self.test_data_dir.joinpath(name).joinpath( - "simulation.egg") - file = he.LocustP3File(str(file_name)) + file = LocustP3File(str(name / egg_filename)) n_streams = file.n_streams n_ch = file.n_channels @@ -159,18 +158,21 @@ def _quick_load_ts_stream(self, name) -> None: ax.set_xlabel(r"Time $[s]$") ax.set_ylabel(r"DAQ V $[V]$") ax.legend(loc="best") - self._save_fig(fig, title) + #self._save_fig(fig, f'{name}/plot_quick_{s}.pdf', title) + plt.show() def test_fft_stream(self) -> None: - for k in self.test_data_dict.keys(): - self._load_fft_stream(k) - self._quick_load_fft_stream(k) + for _, path in self.dataset: + print(path) + self._load_fft_stream(path) + self._quick_load_fft_stream(path) + + ok = input('Plots look ok? (y/n): ').lower().strip() == 'y' + self.assertTrue(ok) def _load_fft_stream(self, name) -> None: - # Plot some specific test data - file_name = self.test_data_dir.joinpath(name).joinpath( - "simulation.egg") - file = he.LocustP3File(str(file_name)) + + file = LocustP3File(str(name / egg_filename)) n_streams = file.n_streams for s in range(n_streams): @@ -192,13 +194,12 @@ def _load_fft_stream(self, name) -> None: ax.set_xlabel(r"Frequency $[Hz]$") ax.set_ylabel(r"FFT") ax.legend(loc="best") - self._save_fig(fig, title) + #self._save_fig(fig, title) + plt.show() def _quick_load_fft_stream(self, name) -> None: # Plot some specific test data - file_name = self.test_data_dir.joinpath(name).joinpath( - "simulation.egg") - file = he.LocustP3File(str(file_name)) + file = LocustP3File(str(name / egg_filename)) n_streams = file.n_streams n_ch = file.n_channels @@ -225,9 +226,10 @@ def _quick_load_fft_stream(self, name) -> None: ax.set_xlabel(r"Frequency $[Hz]$") ax.set_ylabel(r"FFT") ax.legend(loc="best") - self._save_fig(fig, title) + #self._save_fig(fig, title) + plt.show() - def _save_fig(self, fig: Figure, title: str = None) -> None: + def _save_fig(self, fig: Figure, out_file: str, title: str = None) -> None: # Saves the figure to data/images with title and closes the figure. # Default title to axes title if there's only one axes if title is None: @@ -236,11 +238,8 @@ def _save_fig(self, fig: Figure, title: str = None) -> None: title = ax_list[0].get_title() else: raise ValueError("title cannot be None") - out_file = os.path.join( - str(self.test_data_dir), - title.replace(' ', '_').replace('.', '-') + ".pdf") fig.savefig(out_file, bbox_inches='tight', dpi=300) - plt.close(fig) + plt.close(fig) if __name__ == '__main__': diff --git a/test/test_locustP2.py b/test/test_locustP2.py index 65bb3dd..4e03ea0 100644 --- a/test/test_locustP2.py +++ b/test/test_locustP2.py @@ -6,12 +6,15 @@ """ -import hercules as he +from hercules import KassLocustP3, ConfigList, SimConfig from pathlib import Path import numpy as np import unittest +import shutil module_dir = Path(__file__).parent.absolute() +test_dataset_name = 'working_directory' +test_path = module_dir / test_dataset_name class LocustP2Test(unittest.TestCase): def setUp(self) -> None: @@ -20,19 +23,15 @@ def setUp(self) -> None: theta_range = np.linspace(89.7, 90.0, 2) r_phi_range = np.linspace(0, 2 * np.pi, 1) - config_list = [] - sim = he.KassLocustP3(str(module_dir) + '/out_test_locust_P2') + sim = KassLocustP3(test_path, use_kass=True, use_locust=True) + configlist = ConfigList(acquisition_rate=1., info='hello', trap='nonsense') for theta in theta_range: for r_phi in r_phi_range: for r in r_range: x = r * np.cos(r_phi) y = r * np.sin(r_phi) - r_phi_deg = np.rad2deg(r_phi) - name = "Sim_theta_{:.4f}_R_{:.4f}_phi_{:.4f}".format( - theta, r, r_phi_deg) - config = he.SimConfig( - name, + config = SimConfig( phase = 'Phase2', seed_locust=42, seed_kass=43, @@ -49,14 +48,26 @@ def setUp(self) -> None: record_size=3000, v_range=3.0e-6, geometry='Trap_3.xml') - config_list.append(config) + configlist.add_config(config) - sim(config_list) + sim(configlist) + + def tearDown(self) -> None: + shutil.rmtree(test_path) def test_locust(self) -> None: - # Test sim generation (done in setUp), substitutes regular locust test - # This test is always true. - self.assertTrue(True) + + paths =[test_path / 'index.he', + test_path / 'info.txt', + test_path / 'run0' / 'simulation.egg', + test_path / 'run1' / 'simulation.egg', + test_path / 'run2' / 'simulation.egg', + test_path / 'run3' / 'simulation.egg' + ] + + for path in paths: + self.assertTrue(path.exists()) + if __name__ == '__main__': unittest.main()