From 320aab26f66b6eeee0da151926c8ba7b5c589242 Mon Sep 17 00:00:00 2001 From: Eric Sivonxay Date: Thu, 2 Nov 2023 14:10:18 -0700 Subject: [PATCH] Parallel Differential Equation runner --- src/NanoParticleTools/builder.py | 3 +- .../differential_kinetics/__init__.py | 4 ++ .../differential_kinetics/runner.py | 60 +++++++++++++++++++ .../util.py} | 47 +++++++++++++-- tests/differential_kinetics/test_runner.py | 19 ++++++ .../test_util.py} | 25 +++++--- 6 files changed, 141 insertions(+), 17 deletions(-) create mode 100644 src/NanoParticleTools/differential_kinetics/__init__.py create mode 100644 src/NanoParticleTools/differential_kinetics/runner.py rename src/NanoParticleTools/{differential_kinetics.py => differential_kinetics/util.py} (80%) create mode 100644 tests/differential_kinetics/test_runner.py rename tests/{test_differential_kinetics.py => differential_kinetics/test_util.py} (80%) diff --git a/src/NanoParticleTools/builder.py b/src/NanoParticleTools/builder.py index c07b602..acec558 100644 --- a/src/NanoParticleTools/builder.py +++ b/src/NanoParticleTools/builder.py @@ -3,8 +3,7 @@ get_spectrum_wavelength_from_dndt, average_dndt, intensities_from_docs) -from maggma.core import Builder -from maggma.core import Store +from maggma.core import Builder, Store from maggma.utils import grouper from typing import Iterator, List, Dict, Optional, Iterable, Tuple from bson import uuid diff --git a/src/NanoParticleTools/differential_kinetics/__init__.py b/src/NanoParticleTools/differential_kinetics/__init__.py new file mode 100644 index 0000000..e01408f --- /dev/null +++ b/src/NanoParticleTools/differential_kinetics/__init__.py @@ -0,0 +1,4 @@ +from NanoParticleTools.differential_kinetics.util import ( + get_templates, run_one_rate_eq, save_data_to_hdf5, load_data_from_hdf5, + run_and_save_one, get_diff_kinetics_parser) +from NanoParticleTools.differential_kinetics.runner import DifferentialKinetics diff --git a/src/NanoParticleTools/differential_kinetics/runner.py b/src/NanoParticleTools/differential_kinetics/runner.py new file mode 100644 index 0000000..a66f919 --- /dev/null +++ b/src/NanoParticleTools/differential_kinetics/runner.py @@ -0,0 +1,60 @@ +from NanoParticleTools.differential_kinetics import (get_templates, + run_one_rate_eq, + save_data_to_hdf5) +from NanoParticleTools.species_data import Dopant +from maggma.core import Builder, Store +from argparse import ArgumentParser +from h5py import File + +from typing import Iterator + + +class DifferentialKinetics(Builder): + """ + Builder that processes and averages NPMC documents + """ + + def __init__(self, args: ArgumentParser, **kwargs): + + self.source = None + self.args = args + self.target = None + self.kwargs = kwargs + self._file = None + + super().__init__(sources=None, targets=None, chunk_size=1, **kwargs) + + def connect(self): + # Since we aren't using stores, do nothing + return + + @property + def file(self): + if self._file is None: + self._file = File(self.args.output_file, 'w') + return self._file + + def get_items(self) -> Iterator[dict]: + for sample_id, template in enumerate(get_templates(self.args)): + yield (sample_id, template) + + def process_item(self, item: tuple[int, dict]) -> dict: + sample_id, template = item + dopants = [ + Dopant(el, x) + for el, x in zip(template['dopants'], template['dopant_concs']) + ] + output = run_one_rate_eq( + dopants, + excitation_wavelength=template['excitation_wavelength'], + excitation_power=template['excitation_power'], + include_spectra=self.args.include_spectra) + + group_id = int(sample_id // self.args.max_data_per_group) + data_id = int(sample_id % self.args.max_data_per_group) + return (group_id, data_id, output) + + def update_targets(self, items: list[dict]) -> None: + for item in items: + group_id, data_id, output = item + save_data_to_hdf5(self.file, group_id, data_id, output) diff --git a/src/NanoParticleTools/differential_kinetics.py b/src/NanoParticleTools/differential_kinetics/util.py similarity index 80% rename from src/NanoParticleTools/differential_kinetics.py rename to src/NanoParticleTools/differential_kinetics/util.py index a32ab4e..4c7a5a8 100644 --- a/src/NanoParticleTools/differential_kinetics.py +++ b/src/NanoParticleTools/differential_kinetics/util.py @@ -14,8 +14,10 @@ import json import h5py +import logging -def get_parser(): + +def get_diff_kinetics_parser(): parser = argparse.ArgumentParser() parser.add_argument('-n', @@ -85,6 +87,35 @@ def get_parser(): return parser +def get_templates(args): + for sample_id in range(args.num_samples): + # Pick a number of dopants + n_dopants = np.random.choice(range(1, args.max_dopants + 1)) + + # Pick the dopants + dopants = np.random.choice(args.dopants, n_dopants, replace=False) + + # Get the dopant concentrations, normalizing the total concentration to 0-1 + total_conc = np.random.uniform(0, 1) + dopant_concs = np.random.uniform(0, 1, n_dopants) + dopant_concs = total_conc * dopant_concs / np.sum(dopant_concs) + + # sample a wavelength + wavelength = np.random.uniform(*args.excitation_wavelength) + + # sample a power + power = np.random.uniform(*args.excitation_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, + 'excitation_power': power + } + + def run_one_rate_eq(dopants: list[Dopant], excitation_wavelength: int, excitation_power: float, @@ -128,9 +159,10 @@ def run_one_rate_eq(dopants: list[Dopant], 'simulation_time': t[-1], 'excitation_wavelength': excitation_wavelength, 'excitation_power': excitation_power, - 'dopant_concentrations': - {dopant.symbol: dopant.molar_concentration - for dopant in dopants} + 'dopant_concentrations': { + dopant.symbol: dopant.molar_concentration + for dopant in dopants + } } } out_dict['populations'] = final_pop @@ -152,11 +184,14 @@ def run_one_rate_eq(dopants: list[Dopant], def run_and_save_one(dopants: list[str], dopant_concs: list[float], - group_id: int, data_i: int, file: h5py.File, **kwargs): + group_id: int, data_id: int, sample_id: int, + file: h5py.File, **kwargs): dopants = [Dopant(el, x) for el, x in zip(dopants, dopant_concs)] + + logging.info(f'Running Sample {sample_id}') out_dict = run_one_rate_eq(dopants, **kwargs) - save_data_to_hdf5(file, group_id, data_i, out_dict) + save_data_to_hdf5(file, group_id, data_id, out_dict) return out_dict diff --git a/tests/differential_kinetics/test_runner.py b/tests/differential_kinetics/test_runner.py new file mode 100644 index 0000000..6012eb1 --- /dev/null +++ b/tests/differential_kinetics/test_runner.py @@ -0,0 +1,19 @@ +from NanoParticleTools.differential_kinetics.runner import DifferentialKinetics +from NanoParticleTools.differential_kinetics.util import ( + get_diff_kinetics_parser, load_data_from_hdf5) +from h5py import File + + +def test_runner(tmp_path): + args = get_diff_kinetics_parser().parse_args([ + '-n', '2', '-w', '400', '700', '-p', '10', '100000', '-o', + f'{tmp_path}/out.h5', '-s' + ]) + dk = DifferentialKinetics(args) + dk.run() + + data_points = [] + with File(tmp_path / 'out.h5', 'r') as f: + for i in range(0, len(f['group_0'])): + data_points.append(load_data_from_hdf5(f, 0, i)) + assert len(data_points) == 2 diff --git a/tests/test_differential_kinetics.py b/tests/differential_kinetics/test_util.py similarity index 80% rename from tests/test_differential_kinetics.py rename to tests/differential_kinetics/test_util.py index 0b2d357..7b1d442 100644 --- a/tests/test_differential_kinetics.py +++ b/tests/differential_kinetics/test_util.py @@ -1,8 +1,6 @@ -from NanoParticleTools.differential_kinetics import (get_parser, - run_one_rate_eq, - run_and_save_one, - save_data_to_hdf5, - load_data_from_hdf5) +from NanoParticleTools.differential_kinetics.util import ( + get_diff_kinetics_parser, run_one_rate_eq, run_and_save_one, + save_data_to_hdf5, load_data_from_hdf5, get_templates) from NanoParticleTools.species_data import Dopant import numpy as np @@ -17,7 +15,8 @@ def hdf5_file(tmp_path): def test_parser(): - args = get_parser().parse_args(['-n', '10', '-m', '2', '-o', 'test.h5']) + args = get_diff_kinetics_parser().parse_args( + ['-n', '10', '-m', '2', '-o', 'test.h5']) assert args.num_samples == 10 assert args.max_dopants == 2 assert args.output_file == 'test.h5' @@ -27,6 +26,13 @@ def test_parser(): assert args.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) + assert len(list(templates)) == 10 + + def test_run_one_rate_eq(): dopants = [Dopant('Yb', 0.1), Dopant('Er', 0.1)] out_dict = run_one_rate_eq(dopants, @@ -44,9 +50,10 @@ def test_run_and_save_one(hdf5_file): dopant_concs = [0.5, 0.2] run_and_save_one(dopants, dopant_concs, - 0, - 3, - hdf5_file, + group_id=0, + data_id=3, + sample_id=3, + file=hdf5_file, excitation_wavelength=900, excitation_power=1e5, include_spectra=True)