From 082d821c0074b0ff74cde154140347197cc57278 Mon Sep 17 00:00:00 2001 From: Eric Sivonxay Date: Thu, 2 Nov 2023 17:08:54 -0700 Subject: [PATCH] Change to the diff eq runner to allow dumpfn usage --- .../differential_kinetics/runner.py | 44 ++++++++++++++++--- .../differential_kinetics/util.py | 25 ++++++----- .../inputs/spectral_kinetics.py | 10 ++++- tests/differential_kinetics/test_runner.py | 12 ++++- tests/differential_kinetics/test_util.py | 10 ++++- 5 files changed, 79 insertions(+), 22 deletions(-) diff --git a/src/NanoParticleTools/differential_kinetics/runner.py b/src/NanoParticleTools/differential_kinetics/runner.py index a66f919..f43de0e 100644 --- a/src/NanoParticleTools/differential_kinetics/runner.py +++ b/src/NanoParticleTools/differential_kinetics/runner.py @@ -14,10 +14,34 @@ class DifferentialKinetics(Builder): Builder that processes and averages NPMC documents """ - def __init__(self, args: ArgumentParser, **kwargs): + def __init__(self, + num_samples: int, + excitation_wavelength: list[float] = None, + excitation_power: list[float] = None, + possible_dopants: list[str] = None, + max_dopants: int = 4, + include_spectra: bool = True, + output_file: str = 'out.h5', + num_workers: int = 1, + max_data_per_group: int = 100000, + **kwargs): + if excitation_wavelength is None: + excitation_wavelength = [500.0, 1500.0] + if excitation_power is None: + excitation_power = [10.0, 100000.0] + if possible_dopants is None: + possible_dopants = ['Yb', 'Er', 'Tm', 'Nd', 'Ho', 'Eu', 'Sm', 'Dy'] + self.num_samples = num_samples + self.excitation_wavelength = excitation_wavelength + self.excitation_power = excitation_power + self.possible_dopants = possible_dopants + self.max_dopants = max_dopants + self.include_spectra = include_spectra + self.output_file = output_file + self.num_workers = num_workers + self.max_data_per_group = max_data_per_group self.source = None - self.args = args self.target = None self.kwargs = kwargs self._file = None @@ -31,11 +55,17 @@ def connect(self): @property def file(self): if self._file is None: - self._file = File(self.args.output_file, 'w') + self._file = File(self.output_file, 'w') return self._file def get_items(self) -> Iterator[dict]: - for sample_id, template in enumerate(get_templates(self.args)): + templates = get_templates( + num_samples=self.num_samples, + excitation_wavelength=self.excitation_wavelength, + excitation_power=self.excitation_power, + possible_dopants=self.possible_dopants, + max_dopants=self.max_dopants) + for sample_id, template in enumerate(templates): yield (sample_id, template) def process_item(self, item: tuple[int, dict]) -> dict: @@ -48,10 +78,10 @@ def process_item(self, item: tuple[int, dict]) -> dict: dopants, excitation_wavelength=template['excitation_wavelength'], excitation_power=template['excitation_power'], - include_spectra=self.args.include_spectra) + include_spectra=self.include_spectra) - group_id = int(sample_id // self.args.max_data_per_group) - data_id = int(sample_id % self.args.max_data_per_group) + group_id = int(sample_id // self.max_data_per_group) + data_id = int(sample_id % self.max_data_per_group) return (group_id, data_id, output) def update_targets(self, items: list[dict]) -> None: diff --git a/src/NanoParticleTools/differential_kinetics/util.py b/src/NanoParticleTools/differential_kinetics/util.py index 4c7a5a8..3add994 100644 --- a/src/NanoParticleTools/differential_kinetics/util.py +++ b/src/NanoParticleTools/differential_kinetics/util.py @@ -45,7 +45,7 @@ def get_diff_kinetics_parser(): # add optional arguments parser.add_argument( '-d', - '--dopants', + '--possible_dopants', help='The possible dopants to include in the simulation', nargs='+', type=list, @@ -87,13 +87,19 @@ def get_diff_kinetics_parser(): return parser -def get_templates(args): - for sample_id in range(args.num_samples): +def get_templates( + num_samples: int = 4, + excitation_wavelength: list[float] = None, + excitation_power: list[float] = None, + possible_dopants: list[str] = None, + max_dopants: int = 4, +): + for _ in range(num_samples): # Pick a number of dopants - n_dopants = np.random.choice(range(1, args.max_dopants + 1)) + n_dopants = np.random.choice(range(1, max_dopants + 1)) # Pick the dopants - dopants = np.random.choice(args.dopants, n_dopants, replace=False) + dopants = np.random.choice(possible_dopants, n_dopants, replace=False) # Get the dopant concentrations, normalizing the total concentration to 0-1 total_conc = np.random.uniform(0, 1) @@ -101,14 +107,13 @@ def get_templates(args): dopant_concs = total_conc * dopant_concs / np.sum(dopant_concs) # sample a wavelength - wavelength = np.random.uniform(*args.excitation_wavelength) + wavelength = np.random.uniform(*excitation_wavelength) # sample a power - power = np.random.uniform(*args.excitation_power) + power_limits = np.log10(excitation_power) + power = np.random.uniform(*power_limits) + power = np.power(10, power) yield { - # 'sample_id': sample_id, - # 'group_id': sample_id // args.max_data_per_group, - # 'data_id': sample_id % args.max_data_per_group, 'dopants': dopants, 'dopant_concs': dopant_concs, 'excitation_wavelength': wavelength, diff --git a/src/NanoParticleTools/inputs/spectral_kinetics.py b/src/NanoParticleTools/inputs/spectral_kinetics.py index e03efe7..1228695 100644 --- a/src/NanoParticleTools/inputs/spectral_kinetics.py +++ b/src/NanoParticleTools/inputs/spectral_kinetics.py @@ -14,6 +14,9 @@ from scipy.integrate import solve_ivp from functools import lru_cache from monty.json import MSONable +import time + +import logging class SpectralKinetics(MSONable): @@ -109,6 +112,7 @@ def __init__(self, raise ValueError( 'Error, the sum of the dopant concentrations must be <= 1') self.dopants = dopants + self.logger = logging.getLogger(type(self).__name__) @property def mpr_gamma(self) -> float: @@ -597,7 +601,7 @@ def run_kinetics( if isinstance(initial_populations, (list, np.ndarray)): # check if supplied populations is the proper length if len(initial_populations) == self.total_n_levels: - print('Using user input initial population') + self.logger.info('Using user input initial population') else: raise ValueError( "Supplied Population is invalid. Expected length of {self.total_n_levels, " @@ -617,10 +621,14 @@ def dNdt_fn(t, y): if 'method' not in kwargs: kwargs['method'] = 'BDF' + + start_time = time.time() sol = solve_ivp(get_dNdt_fn(self), t_span=t_span, y0=initial_populations, **kwargs) + end_time = time.time() + self.logger.info(f'Found solution in {end_time-start_time:.2f} seconds') return sol.t, sol.y.T diff --git a/tests/differential_kinetics/test_runner.py b/tests/differential_kinetics/test_runner.py index 6012eb1..05e4ff2 100644 --- a/tests/differential_kinetics/test_runner.py +++ b/tests/differential_kinetics/test_runner.py @@ -6,10 +6,12 @@ def test_runner(tmp_path): args = get_diff_kinetics_parser().parse_args([ - '-n', '2', '-w', '400', '700', '-p', '10', '100000', '-o', + '-n', '2', '-m', '2', '-w', '400', '700', '-p', '10', '100000', '-o', f'{tmp_path}/out.h5', '-s' ]) - dk = DifferentialKinetics(args) + dk = DifferentialKinetics(**vars(args)) + assert dk.num_samples == 2 + dk.run() data_points = [] @@ -17,3 +19,9 @@ def test_runner(tmp_path): for i in range(0, len(f['group_0'])): data_points.append(load_data_from_hdf5(f, 0, i)) assert len(data_points) == 2 + + dk = DifferentialKinetics(1000) + assert dk.num_samples == 1000 + assert dk.possible_dopants == ['Yb', 'Er', 'Tm', 'Nd', 'Ho', 'Eu', 'Sm', 'Dy'] + assert dk.excitation_wavelength == [500.0, 1500.0] + assert dk.excitation_power == [10.0, 100000.0] diff --git a/tests/differential_kinetics/test_util.py b/tests/differential_kinetics/test_util.py index 7b1d442..d8749d8 100644 --- a/tests/differential_kinetics/test_util.py +++ b/tests/differential_kinetics/test_util.py @@ -23,13 +23,19 @@ def test_parser(): assert args.include_spectra is False assert args.excitation_wavelength == (700, 1500) assert args.excitation_power == (10, 1e6) - assert args.dopants == ["Yb", "Er", "Tm", "Nd", "Ho", "Eu", "Sm", "Dy"] + assert args.possible_dopants == [ + "Yb", "Er", "Tm", "Nd", "Ho", "Eu", "Sm", "Dy" + ] def test_get_templates(): args = get_diff_kinetics_parser().parse_args( ['-n', '10', '-m', '2', '-o', 'test.h5']) - templates = get_templates(args) + templates = get_templates(num_samples=args.num_samples, + excitation_wavelength=args.excitation_wavelength, + max_dopants=args.max_dopants, + excitation_power=args.excitation_power, + possible_dopants=args.possible_dopants) assert len(list(templates)) == 10