From 55ebbd7e43ba4d9ebd0958e9170ccca675546e4b Mon Sep 17 00:00:00 2001 From: Dave Grote Date: Wed, 16 Mar 2022 10:27:29 -0700 Subject: [PATCH 01/13] Added test of type checking and autoargs --- PICMI_Python/applied_fields.py | 31 +++++++++++++++++-------------- PICMI_Python/base.py | 13 +++++++++++++ PICMI_Python/requirements.txt | 2 ++ 3 files changed, 32 insertions(+), 14 deletions(-) diff --git a/PICMI_Python/applied_fields.py b/PICMI_Python/applied_fields.py index 74def08..054b4a7 100644 --- a/PICMI_Python/applied_fields.py +++ b/PICMI_Python/applied_fields.py @@ -2,14 +2,20 @@ These should be the base classes for Python implementation of the PICMI standard """ import re +import typing +from collections.abc import Sequence -from .base import _ClassWithInit +from autoclass import autoargs +from typeguard import typechecked + +from .base import _ClassWithInit, VectorFloat3 # --------------- # Applied fields # --------------- +@typechecked class PICMI_ConstantAppliedField(_ClassWithInit): """ Describes a constant applied field @@ -22,19 +28,16 @@ class PICMI_ConstantAppliedField(_ClassWithInit): - lower_bound=[None,None,None]: Lower bound of the region where the field is applied (vector) [m] - upper_bound=[None,None,None]: Upper bound of the region where the field is applied (vector) [m] """ - def __init__(self, Ex=None, Ey=None, Ez=None, Bx=None, By=None, Bz=None, - lower_bound=[None,None,None], upper_bound=[None,None,None], - **kw): - - self.Ex = Ex - self.Ey = Ey - self.Ez = Ez - self.Bx = Bx - self.By = By - self.Bz = Bz - - self.lower_bound = lower_bound - self.upper_bound = upper_bound + @autoargs(exclude=['kw']) + def __init__(self, Ex : float = None, + Ey : float = None, + Ez : float = None, + Bx : float = None, + By : float = None, + Bz : float = None, + lower_bound : VectorFloat3 = [None,None,None], + upper_bound : VectorFloat3 = [None,None,None], + **kw): self.handle_init(kw) diff --git a/PICMI_Python/base.py b/PICMI_Python/base.py index e6cb526..3cd1139 100644 --- a/PICMI_Python/base.py +++ b/PICMI_Python/base.py @@ -2,6 +2,9 @@ """ import inspect import warnings +import typing +from collections.abc import Sequence + codename = None @@ -26,8 +29,18 @@ def register_constants(implementation_constants): def _get_constants(): return _implementation_constants +VectorFloat3 = typing.NewType('VectorFloat3', Sequence[float]) +VectorInt3 = typing.NewType('VectorFloat3', Sequence[int]) + class _ClassWithInit(object): + def _check_vector_lengths(self): + for arg_name, arg_type in self.__init__.__annotations__.items(): + if arg_type in [VectorFloat3, VectorInt3]: + assert len(getattr(self, arg_name)) == 3, Exception(f'{arg_name} must have a length of 3') + def handle_init(self, kw): + self._check_vector_lengths() + # --- Grab all keywords for the current code. # --- Arguments for other supported codes are ignored. # --- If there is anything left over, it is an error. diff --git a/PICMI_Python/requirements.txt b/PICMI_Python/requirements.txt index c0ad38b..a798097 100644 --- a/PICMI_Python/requirements.txt +++ b/PICMI_Python/requirements.txt @@ -1,2 +1,4 @@ numpy~=1.15 scipy~=1.5 +autoclass +typeguard From eadeb86a0db0fd2db7f9286e88ed28a60bceedee Mon Sep 17 00:00:00 2001 From: Dave Grote Date: Wed, 16 Mar 2022 12:11:25 -0700 Subject: [PATCH 02/13] More changes --- PICMI_Python/applied_fields.py | 56 +++++++++++----------------------- PICMI_Python/base.py | 16 +++++++++- 2 files changed, 33 insertions(+), 39 deletions(-) diff --git a/PICMI_Python/applied_fields.py b/PICMI_Python/applied_fields.py index 054b4a7..21e05ef 100644 --- a/PICMI_Python/applied_fields.py +++ b/PICMI_Python/applied_fields.py @@ -8,7 +8,7 @@ from autoclass import autoargs from typeguard import typechecked -from .base import _ClassWithInit, VectorFloat3 +from .base import _ClassWithInit, VectorFloat3, Expression # --------------- # Applied fields @@ -42,6 +42,7 @@ def __init__(self, Ex : float = None, self.handle_init(kw) +@typechecked class PICMI_AnalyticAppliedField(_ClassWithInit): """ Describes an analytic applied field @@ -57,35 +58,16 @@ class PICMI_AnalyticAppliedField(_ClassWithInit): - lower_bound=[None,None,None]: Lower bound of the region where the field is applied (vector) [m] - upper_bound=[None,None,None]: Upper bound of the region where the field is applied (vector) [m] """ - def __init__(self, Ex_expression=None, Ey_expression=None, Ez_expression=None, - Bx_expression=None, By_expression=None, Bz_expression=None, - lower_bound=[None,None,None], upper_bound=[None,None,None], - **kw): - - self.Ex_expression = Ex_expression - self.Ey_expression = Ey_expression - self.Ez_expression = Ez_expression - self.Bx_expression = Bx_expression - self.By_expression = By_expression - self.Bz_expression = Bz_expression - - self.lower_bound = lower_bound - self.upper_bound = upper_bound - - # --- Find any user defined keywords in the kw dictionary. - # --- Save them and delete them from kw. - # --- It's up to the code to make sure that all parameters - # --- used in the expression are defined. - self.user_defined_kw = {} - for k in list(kw.keys()): - if ((self.Ex_expression is not None and re.search(r'\b%s\b'%k, self.Ex_expression)) or - (self.Ey_expression is not None and re.search(r'\b%s\b'%k, self.Ey_expression)) or - (self.Ez_expression is not None and re.search(r'\b%s\b'%k, self.Ez_expression)) or - (self.Bx_expression is not None and re.search(r'\b%s\b'%k, self.Bx_expression)) or - (self.By_expression is not None and re.search(r'\b%s\b'%k, self.By_expression)) or - (self.Bz_expression is not None and re.search(r'\b%s\b'%k, self.Bz_expression))): - self.user_defined_kw[k] = kw[k] - del kw[k] + @autoargs(exclude=['kw']) + def __init__(self, Ex_expression : Expression = None, + Ey_expression : Expression = None, + Ez_expression : Expression = None, + Bx_expression : Expression = None, + By_expression : Expression = None, + Bz_expression : Expression = None, + lower_bound : VectorFloat3 = [None,None,None], + upper_bound : VectorFloat3 = [None,None,None], + **kw): self.handle_init(kw) @@ -106,17 +88,15 @@ class PICMI_Mirror(_ClassWithInit): or the code's default value if neither are specified. """ - def __init__(self, x_front_location=None, y_front_location=None, z_front_location=None, - depth=None, number_of_cells=None, **kw): + def __init__(self, x_front_location : float = None, + y_front_location : float = None, + z_front_location : float = None, + depth : float = None, + number_of_cells : int = None, + **kw): assert [x_front_location,y_front_location,z_front_location].count(None) == 2,\ Exception('At least one and only one of [x,y,z]_front_location should be specified.') - self.x_front_location = x_front_location - self.y_front_location = y_front_location - self.z_front_location = z_front_location - self.depth = depth - self.number_of_cells = number_of_cells - self.handle_init(kw) diff --git a/PICMI_Python/base.py b/PICMI_Python/base.py index 3cd1139..719eeac 100644 --- a/PICMI_Python/base.py +++ b/PICMI_Python/base.py @@ -4,6 +4,7 @@ import warnings import typing from collections.abc import Sequence +import re codename = None @@ -31,15 +32,28 @@ def _get_constants(): VectorFloat3 = typing.NewType('VectorFloat3', Sequence[float]) VectorInt3 = typing.NewType('VectorFloat3', Sequence[int]) +Expression = typing.NewType('Expression', str) class _ClassWithInit(object): def _check_vector_lengths(self): for arg_name, arg_type in self.__init__.__annotations__.items(): if arg_type in [VectorFloat3, VectorInt3]: - assert len(getattr(self, arg_name)) == 3, Exception(f'{arg_name} must have a length of 3') + arg_value = getattr(self, arg_name) + assert len(arg_value) == 3, Exception(f'{arg_name} must have a length of 3') + + def _check_expressions(self, kw): + for arg_name, arg_type in self.__init__.__annotations__.items(): + if arg_type == Expression: + self.user_defined_kw = getattr(self, 'user_defined_kw', {}) + arg_value = getattr(self, arg_name) + if arg_value is not None: + for k in list(kw.keys()): + if re.search(r'\b%s\b'%k, arg_value): + self.user_defined_kw[k] = kw.pop(k) def handle_init(self, kw): self._check_vector_lengths() + self._check_expressions(kw) # --- Grab all keywords for the current code. # --- Arguments for other supported codes are ignored. From e39922a30152cb361ae9a56d82b4c3a053b6c8f5 Mon Sep 17 00:00:00 2001 From: Dave Grote Date: Wed, 16 Mar 2022 13:32:45 -0700 Subject: [PATCH 03/13] Several fixes and clean up --- PICMI_Python/base.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/PICMI_Python/base.py b/PICMI_Python/base.py index 719eeac..3bf4060 100644 --- a/PICMI_Python/base.py +++ b/PICMI_Python/base.py @@ -30,10 +30,12 @@ def register_constants(implementation_constants): def _get_constants(): return _implementation_constants + VectorFloat3 = typing.NewType('VectorFloat3', Sequence[float]) -VectorInt3 = typing.NewType('VectorFloat3', Sequence[int]) +VectorInt3 = typing.NewType('VectorInt3', Sequence[int]) Expression = typing.NewType('Expression', str) + class _ClassWithInit(object): def _check_vector_lengths(self): for arg_name, arg_type in self.__init__.__annotations__.items(): @@ -41,9 +43,13 @@ def _check_vector_lengths(self): arg_value = getattr(self, arg_name) assert len(arg_value) == 3, Exception(f'{arg_name} must have a length of 3') - def _check_expressions(self, kw): + def _process_expression_arguments(self, kw): + """For arguments that are of type Expression, save any keyword arguments used in + the expression in the user_defined_kw dictionary. + """ for arg_name, arg_type in self.__init__.__annotations__.items(): if arg_type == Expression: + # Create the dictionary is needed self.user_defined_kw = getattr(self, 'user_defined_kw', {}) arg_value = getattr(self, arg_name) if arg_value is not None: @@ -53,7 +59,7 @@ def _check_expressions(self, kw): def handle_init(self, kw): self._check_vector_lengths() - self._check_expressions(kw) + self._process_expression_arguments(kw) # --- Grab all keywords for the current code. # --- Arguments for other supported codes are ignored. From ab6cf4f327b5bf602d783ea8e1528b14ea9ee84e Mon Sep 17 00:00:00 2001 From: Dave Grote Date: Thu, 17 Mar 2022 15:54:29 -0700 Subject: [PATCH 04/13] Allowed numpy arrays for vector arguments --- PICMI_Python/base.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/PICMI_Python/base.py b/PICMI_Python/base.py index 3bf4060..30706cf 100644 --- a/PICMI_Python/base.py +++ b/PICMI_Python/base.py @@ -6,6 +6,9 @@ from collections.abc import Sequence import re +import numpy as np +import numpy.typing as npt + codename = None @@ -31,8 +34,8 @@ def _get_constants(): return _implementation_constants -VectorFloat3 = typing.NewType('VectorFloat3', Sequence[float]) -VectorInt3 = typing.NewType('VectorInt3', Sequence[int]) +VectorFloat3 = typing.NewType('VectorFloat3', typing.Union[Sequence[float], npt.NDArray[np.float64]]) +VectorInt3 = typing.NewType('VectorInt3', typing.Union[Sequence[int], npt.NDArray[np.int64]]) Expression = typing.NewType('Expression', str) From 83773753e6fcbb714123e1c82d707eff73e584e2 Mon Sep 17 00:00:00 2001 From: Dave Grote Date: Thu, 17 Mar 2022 16:24:57 -0700 Subject: [PATCH 05/13] Use np.ndarray instead of npt.NDArray (which requires numpy version 1.20) --- PICMI_Python/base.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/PICMI_Python/base.py b/PICMI_Python/base.py index 30706cf..7417d05 100644 --- a/PICMI_Python/base.py +++ b/PICMI_Python/base.py @@ -7,7 +7,6 @@ import re import numpy as np -import numpy.typing as npt codename = None @@ -34,8 +33,8 @@ def _get_constants(): return _implementation_constants -VectorFloat3 = typing.NewType('VectorFloat3', typing.Union[Sequence[float], npt.NDArray[np.float64]]) -VectorInt3 = typing.NewType('VectorInt3', typing.Union[Sequence[int], npt.NDArray[np.int64]]) +VectorFloat3 = typing.NewType('VectorFloat3', typing.Union[Sequence[float], np.ndarray[3, np.float64]]) +VectorInt3 = typing.NewType('VectorInt3', typing.Union[Sequence[int], np.ndarray[3, np.int64]]) Expression = typing.NewType('Expression', str) From b4b9423169d7c2155a24811a10c95cd08547bd04 Mon Sep 17 00:00:00 2001 From: Dave Grote Date: Fri, 18 Mar 2022 09:20:34 -0700 Subject: [PATCH 06/13] Updated diagnstics.py --- PICMI_Python/__init__.py | 4 +- PICMI_Python/applied_fields.py | 1 - PICMI_Python/diagnostics.py | 208 ++++++++++++++++----------------- 3 files changed, 102 insertions(+), 111 deletions(-) diff --git a/PICMI_Python/__init__.py b/PICMI_Python/__init__.py index 509e935..b39bc24 100644 --- a/PICMI_Python/__init__.py +++ b/PICMI_Python/__init__.py @@ -1,7 +1,7 @@ from .base import * -from .diagnostics import * from .fields import * +from .particles import * +from .diagnostics import * from .applied_fields import * from .lasers import * -from .particles import * from .simulation import * diff --git a/PICMI_Python/applied_fields.py b/PICMI_Python/applied_fields.py index 21e05ef..b5cddcc 100644 --- a/PICMI_Python/applied_fields.py +++ b/PICMI_Python/applied_fields.py @@ -3,7 +3,6 @@ """ import re import typing -from collections.abc import Sequence from autoclass import autoargs from typeguard import typechecked diff --git a/PICMI_Python/diagnostics.py b/PICMI_Python/diagnostics.py index 8eb9cbf..53a39c0 100644 --- a/PICMI_Python/diagnostics.py +++ b/PICMI_Python/diagnostics.py @@ -2,14 +2,32 @@ These should be the base classes for Python implementation of the PICMI standard The classes in the file are all diagnostics related """ +import typing +from collections.abc import Sequence + +from autoclass import autoargs +from typeguard import typechecked from .base import _ClassWithInit +from . import fields +from . import particles + +GridType = typing.NewType('GridType', typing.Union[fields.PICMI_Cartesian1DGrid, + fields.PICMI_CylindricalGrid, + fields.PICMI_Cartesian2DGrid, + fields.PICMI_Cartesian3DGrid]) + +SpeciesType = typing.NewType('SpeciesType', typing.Union[particles.PICMI_Species, + particles.PICMI_MultiSpecies]) +SpeciesArgument = typing.NewType('SpeciesArgument', typing.Union[SpeciesType, Sequence[SpeciesType]]) + # ---------------------------- # Simulation frame diagnostics # ---------------------------- +@typechecked class PICMI_FieldDiagnostic(_ClassWithInit): """ Defines the electromagnetic field diagnostics in the simulation frame @@ -32,35 +50,34 @@ class PICMI_FieldDiagnostic(_ClassWithInit): - name: Sets the base name for the diagnostic output files (optional) """ - def __init__(self, grid, period, data_list=None, - write_dir = None, - step_min = None, - step_max = None, - number_of_cells = None, - lower_bound = None, - upper_bound = None, - parallelio = None, - name = None, - **kw): - - if data_list is not None: - assert isinstance(data_list, list), 'FieldDiagnostic: data_list must be a list' - - self.grid = grid - self.period = period - self.data_list = data_list - self.write_dir = write_dir - self.step_min = step_min - self.step_max = step_max - self.number_of_cells = number_of_cells - self.lower_bound = lower_bound - self.upper_bound = upper_bound - self.parallelio = parallelio - self.name = name + @autoargs(exclude=['kw']) + def __init__(self, grid : GridType, + period : int, + data_list : Sequence[str] = None, + write_dir : str = None, + step_min : int = None, + step_max : int = None, + number_of_cells : Sequence[int] = None, + lower_bound : Sequence[float] = None, + upper_bound : Sequence[float] = None, + parallelio : bool = None, + name : str = None, + **kw): + + if number_of_cells is not None: + assert len(number_of_cells) == grid.number_of_dimensions, \ + 'FieldDiagnostic: length of number_of_cells must be the same as the dimensionality of the grid' + if lower_bound is not None: + assert len(lower_bound) == grid.number_of_dimensions, \ + 'FieldDiagnostic: length of lower_bound must be the same as the dimensionality of the grid' + if upper_bound is not None: + assert len(upper_bound) == grid.number_of_dimensions, \ + 'FieldDiagnostic: length of upper_bound must be the same as the dimensionality of the grid' self.handle_init(kw) +@typechecked class PICMI_ElectrostaticFieldDiagnostic(_ClassWithInit): """ Defines the electrostatic field diagnostics in the simulation frame @@ -82,35 +99,34 @@ class PICMI_ElectrostaticFieldDiagnostic(_ClassWithInit): - parallelio=None: If set to True, field diagnostics are dumped in parallel (optional) - name: Sets the base name for the diagnostic output files (optional) """ - def __init__(self, grid, period, data_list=None, - write_dir = None, - step_min = None, - step_max = None, - number_of_cells = None, - lower_bound = None, - upper_bound = None, - parallelio = None, - name = None, - **kw): - - if data_list is not None: - assert isinstance(data_list, list), 'ElectrostaticFieldDiagnostic: data_list must be a list' - - self.grid = grid - self.period = period - self.data_list = data_list - self.write_dir = write_dir - self.step_min = step_min - self.step_max = step_max - self.number_of_cells = number_of_cells - self.lower_bound = lower_bound - self.upper_bound = upper_bound - self.parallelio = parallelio - self.name = name + @autoargs(exclude=['kw']) + def __init__(self, grid : GridType, + period : int, + data_list : Sequence[str] = None, + write_dir : str = None, + step_min : int = None, + step_max : int = None, + number_of_cells : Sequence[int] = None, + lower_bound : Sequence[float] = None, + upper_bound : Sequence[float] = None, + parallelio : bool = None, + name : str = None, + **kw): + + if number_of_cells is not None: + assert len(number_of_cells) == grid.number_of_dimensions, \ + 'ElectrostaticFieldDiagnostic: length of number_of_cells must be the same as the dimensionality of the grid' + if lower_bound is not None: + assert len(lower_bound) == grid.number_of_dimensions, \ + 'ElectrostaticFieldDiagnostic: length of lower_bound must be the same as the dimensionality of the grid' + if upper_bound is not None: + assert len(upper_bound) == grid.number_of_dimensions, \ + 'ElectrostaticFieldDiagnostic: length of upper_bound must be the same as the dimensionality of the grid' self.handle_init(kw) +@typechecked class PICMI_ParticleDiagnostic(_ClassWithInit) : """ Defines the particle diagnostics in the simulation frame @@ -128,25 +144,16 @@ class PICMI_ParticleDiagnostic(_ClassWithInit) : - name: Sets the base name for the diagnostic output files (optional) """ - def __init__(self, period, species, data_list=None, - write_dir = None, - step_min = None, - step_max = None, - parallelio = None, - name = None, - **kw): - - if data_list is not None: - assert isinstance(data_list, list), 'ParticleDiagnostic: data_list must be a list' - - self.period = period - self.species = species - self.data_list = data_list - self.write_dir = write_dir - self.step_min = step_min - self.step_max = step_max - self.parallelio = parallelio - self.name = name + @autoargs(exclude=['kw']) + def __init__(self, period : int, + species : SpeciesArgument, + data_list : Sequence[str] = None, + write_dir : str = None, + step_min : int = None, + step_max : int = None, + parallelio : bool = None, + name : str = None, + **kw): self.handle_init(kw) @@ -156,6 +163,7 @@ def __init__(self, period, species, data_list=None, # ---------------------------- +@typechecked class PICMI_LabFrameFieldDiagnostic(_ClassWithInit): """ Defines the electromagnetic field diagnostics in the lab frame @@ -170,29 +178,22 @@ class PICMI_LabFrameFieldDiagnostic(_ClassWithInit): - parallelio=None: If set to True, field diagnostics are dumped in parallel (optional) - name: Sets the base name for the diagnostic output files (optional) """ - def __init__(self, grid, num_snapshots, dt_snapshots, data_list=None, - z_subsampling = 1, time_start = 0., - write_dir = None, - parallelio = None, - name = None, - **kw): - - if data_list is not None: - assert isinstance(data_list, list), 'LabFrameFieldDiagnostic: data_list must be a list' - - self.grid = grid - self.num_snapshots = num_snapshots - self.dt_snapshots = dt_snapshots - self.z_subsampling = z_subsampling - self.time_start = time_start - self.data_list = data_list - self.write_dir = write_dir - self.parallelio = parallelio - self.name = name + @autoargs(exclude=['kw']) + def __init__(self, grid : GridType, + num_snapshots : int, + dt_snapshots : float, + data_list : Sequence[str] = None, + z_subsampling : int = 1, + time_start : float = 0., + write_dir : str = None, + parallelio : bool = None, + name : str = None, + **kw): self.handle_init(kw) +@typechecked class PICMI_LabFrameParticleDiagnostic(_ClassWithInit): """ Defines the particle diagnostics in the lab frame @@ -208,25 +209,16 @@ class PICMI_LabFrameParticleDiagnostic(_ClassWithInit): - parallelio=None: If set to True, particle diagnostics are dumped in parallel (optional) - name: Sets the base name for the diagnostic output files (optional) """ - def __init__(self, grid, num_snapshots, dt_snapshots, data_list=None, - time_start = 0., - species = None, - write_dir = None, - parallelio = None, - name = None, - **kw): - - if data_list is not None: - assert isinstance(data_list, list), 'LabFrameParticleDiagnostic: data_list must be a list' - - self.grid = grid - self.num_snapshots = num_snapshots - self.dt_snapshots = dt_snapshots - self.time_start = time_start - self.species = species - self.data_list = data_list - self.write_dir = write_dir - self.parallelio = parallelio - self.name = name + @autoargs(exclude=['kw']) + def __init__(self, grid : GridType, + num_snapshots : int, + dt_snapshots : float, + data_list : Sequence[str] = None, + time_start : float = 0., + species : SpeciesArgument = None, + write_dir : str = None, + parallelio : bool = None, + name : str = None, + **kw): self.handle_init(kw) From e12ccaf185f462b14ddbba170757b3be3cc386c3 Mon Sep 17 00:00:00 2001 From: Dave Grote Date: Mon, 21 Mar 2022 10:19:28 -0700 Subject: [PATCH 07/13] Moved type definitions into picmi_types.py --- PICMI_Python/applied_fields.py | 23 ++++++++++++----------- PICMI_Python/base.py | 22 +++++++++++++--------- PICMI_Python/diagnostics.py | 22 +++++++--------------- PICMI_Python/picmi_types.py | 24 ++++++++++++++++++++++++ 4 files changed, 56 insertions(+), 35 deletions(-) create mode 100644 PICMI_Python/picmi_types.py diff --git a/PICMI_Python/applied_fields.py b/PICMI_Python/applied_fields.py index b5cddcc..72b33ca 100644 --- a/PICMI_Python/applied_fields.py +++ b/PICMI_Python/applied_fields.py @@ -7,7 +7,8 @@ from autoclass import autoargs from typeguard import typechecked -from .base import _ClassWithInit, VectorFloat3, Expression +from .base import _ClassWithInit +from . import picmi_types # --------------- # Applied fields @@ -34,8 +35,8 @@ def __init__(self, Ex : float = None, Bx : float = None, By : float = None, Bz : float = None, - lower_bound : VectorFloat3 = [None,None,None], - upper_bound : VectorFloat3 = [None,None,None], + lower_bound : picmi_types.VectorFloat3 = [None,None,None], + upper_bound : picmi_types.VectorFloat3 = [None,None,None], **kw): self.handle_init(kw) @@ -58,14 +59,14 @@ class PICMI_AnalyticAppliedField(_ClassWithInit): - upper_bound=[None,None,None]: Upper bound of the region where the field is applied (vector) [m] """ @autoargs(exclude=['kw']) - def __init__(self, Ex_expression : Expression = None, - Ey_expression : Expression = None, - Ez_expression : Expression = None, - Bx_expression : Expression = None, - By_expression : Expression = None, - Bz_expression : Expression = None, - lower_bound : VectorFloat3 = [None,None,None], - upper_bound : VectorFloat3 = [None,None,None], + def __init__(self, Ex_expression : picmi_types.Expression = None, + Ey_expression : picmi_types.Expression = None, + Ez_expression : picmi_types.Expression = None, + Bx_expression : picmi_types.Expression = None, + By_expression : picmi_types.Expression = None, + Bz_expression : picmi_types.Expression = None, + lower_bound : picmi_types.VectorFloat3 = [None,None,None], + upper_bound : picmi_types.VectorFloat3 = [None,None,None], **kw): self.handle_init(kw) diff --git a/PICMI_Python/base.py b/PICMI_Python/base.py index 7417d05..31346b1 100644 --- a/PICMI_Python/base.py +++ b/PICMI_Python/base.py @@ -8,6 +8,8 @@ import numpy as np +from . import picmi_types + codename = None @@ -33,15 +35,10 @@ def _get_constants(): return _implementation_constants -VectorFloat3 = typing.NewType('VectorFloat3', typing.Union[Sequence[float], np.ndarray[3, np.float64]]) -VectorInt3 = typing.NewType('VectorInt3', typing.Union[Sequence[int], np.ndarray[3, np.int64]]) -Expression = typing.NewType('Expression', str) - - class _ClassWithInit(object): def _check_vector_lengths(self): for arg_name, arg_type in self.__init__.__annotations__.items(): - if arg_type in [VectorFloat3, VectorInt3]: + if arg_type in [picmi_types.VectorFloat3, picmi_types.VectorInt3]: arg_value = getattr(self, arg_name) assert len(arg_value) == 3, Exception(f'{arg_name} must have a length of 3') @@ -50,10 +47,17 @@ def _process_expression_arguments(self, kw): the expression in the user_defined_kw dictionary. """ for arg_name, arg_type in self.__init__.__annotations__.items(): - if arg_type == Expression: - # Create the dictionary is needed - self.user_defined_kw = getattr(self, 'user_defined_kw', {}) + if arg_type == picmi_types.Expression: + arg_value = getattr(self, arg_name) + + # --- Remove any line feeds from the expression + if arg_value is not None: + arg_value = arg_value.replace('\n', '') + setattr(self, arg_name, arg_value) + + # --- The dictionary is created if needed + self.user_defined_kw = getattr(self, 'user_defined_kw', {}) if arg_value is not None: for k in list(kw.keys()): if re.search(r'\b%s\b'%k, arg_value): diff --git a/PICMI_Python/diagnostics.py b/PICMI_Python/diagnostics.py index 53a39c0..7542e9c 100644 --- a/PICMI_Python/diagnostics.py +++ b/PICMI_Python/diagnostics.py @@ -11,15 +11,7 @@ from .base import _ClassWithInit from . import fields from . import particles - -GridType = typing.NewType('GridType', typing.Union[fields.PICMI_Cartesian1DGrid, - fields.PICMI_CylindricalGrid, - fields.PICMI_Cartesian2DGrid, - fields.PICMI_Cartesian3DGrid]) - -SpeciesType = typing.NewType('SpeciesType', typing.Union[particles.PICMI_Species, - particles.PICMI_MultiSpecies]) -SpeciesArgument = typing.NewType('SpeciesArgument', typing.Union[SpeciesType, Sequence[SpeciesType]]) +from . import picmi_types # ---------------------------- @@ -51,7 +43,7 @@ class PICMI_FieldDiagnostic(_ClassWithInit): """ @autoargs(exclude=['kw']) - def __init__(self, grid : GridType, + def __init__(self, grid : picmi_types.GridType, period : int, data_list : Sequence[str] = None, write_dir : str = None, @@ -100,7 +92,7 @@ class PICMI_ElectrostaticFieldDiagnostic(_ClassWithInit): - name: Sets the base name for the diagnostic output files (optional) """ @autoargs(exclude=['kw']) - def __init__(self, grid : GridType, + def __init__(self, grid : picmi_types.GridType, period : int, data_list : Sequence[str] = None, write_dir : str = None, @@ -146,7 +138,7 @@ class PICMI_ParticleDiagnostic(_ClassWithInit) : @autoargs(exclude=['kw']) def __init__(self, period : int, - species : SpeciesArgument, + species : picmi_types.SpeciesArgument, data_list : Sequence[str] = None, write_dir : str = None, step_min : int = None, @@ -179,7 +171,7 @@ class PICMI_LabFrameFieldDiagnostic(_ClassWithInit): - name: Sets the base name for the diagnostic output files (optional) """ @autoargs(exclude=['kw']) - def __init__(self, grid : GridType, + def __init__(self, grid : picmi_types.GridType, num_snapshots : int, dt_snapshots : float, data_list : Sequence[str] = None, @@ -210,12 +202,12 @@ class PICMI_LabFrameParticleDiagnostic(_ClassWithInit): - name: Sets the base name for the diagnostic output files (optional) """ @autoargs(exclude=['kw']) - def __init__(self, grid : GridType, + def __init__(self, grid : picmi_types.GridType, num_snapshots : int, dt_snapshots : float, data_list : Sequence[str] = None, time_start : float = 0., - species : SpeciesArgument = None, + species : picmi_types.SpeciesArgument = None, write_dir : str = None, parallelio : bool = None, name : str = None, diff --git a/PICMI_Python/picmi_types.py b/PICMI_Python/picmi_types.py new file mode 100644 index 0000000..e5a1936 --- /dev/null +++ b/PICMI_Python/picmi_types.py @@ -0,0 +1,24 @@ +"""Types definitions for the picmi standard +""" +import typing +from collections.abc import Sequence + +import numpy as np + +VectorFloat3 = typing.NewType('VectorFloat3', typing.Union[Sequence[float], np.ndarray[3, np.float64]]) +VectorInt3 = typing.NewType('VectorInt3', typing.Union[Sequence[int], np.ndarray[3, np.int64]]) +Expression = typing.NewType('Expression', str) + + +# These must be defined as strings to avoid circular definitions +GridType = """typing.NewType('GridType', typing.Union[fields.PICMI_Cartesian1DGrid, + fields.PICMI_CylindricalGrid, + fields.PICMI_Cartesian2DGrid, + fields.PICMI_Cartesian3DGrid])""" + +SpeciesType = """typing.NewType('SpeciesType', typing.Union[particles.PICMI_Species, + particles.PICMI_MultiSpecies])""" + +SpeciesArgument = """typing.NewType('SpeciesArgument', typing.Union[SpeciesType, Sequence[SpeciesType]])""" + + From e27b76aeceefcc0029725e0fc38e88087cb6efe5 Mon Sep 17 00:00:00 2001 From: Dave Grote Date: Mon, 21 Mar 2022 10:20:00 -0700 Subject: [PATCH 08/13] Added types to lasers.py --- PICMI_Python/lasers.py | 114 ++++++++++++++++------------------------- 1 file changed, 43 insertions(+), 71 deletions(-) diff --git a/PICMI_Python/lasers.py b/PICMI_Python/lasers.py index bb1ea0e..bf7fbd6 100644 --- a/PICMI_Python/lasers.py +++ b/PICMI_Python/lasers.py @@ -5,12 +5,17 @@ import sys import re +from autoclass import autoargs +from typeguard import typechecked + from .base import _ClassWithInit, _get_constants +from . import picmi_types # --------------- # Physics objects # --------------- +@typechecked class PICMI_GaussianLaser(_ClassWithInit): """ Specifies a Gaussian laser distribution. @@ -51,51 +56,38 @@ class PICMI_GaussianLaser(_ClassWithInit): - beta: Angular dispersion at focus (in the lab frame) [rad.s] - phi2: Temporal chirp at focus (in the lab frame) [s^2] - fill_in=True: Flags whether to fill in the empty spaced opened up when the grid moves - name=None: Optional name of the laser + - name=None: Optional name of the laser """ - def __init__(self, wavelength, waist, duration, - focal_position = [0., 0., 0.], - centroid_position = [0., 0., 0.], - propagation_direction = [0., 0., 1.], - polarization_direction = [1., 0., 0.], - a0 = None, - E0 = None, - phi0 = None, - zeta = None, - beta = None, - phi2 = None, - name = None, - fill_in = True, - **kw): + @autoargs(exclude=['kw']) + def __init__(self, wavelength : float, + waist : float, + duration : float, + focal_position : picmi_types.VectorFloat3 = [0., 0., 0.], + centroid_position : picmi_types.VectorFloat3 = [0., 0., 0.], + propagation_direction : picmi_types.VectorFloat3 = [0., 0., 1.], + polarization_direction : picmi_types.VectorFloat3 = [1., 0., 0.], + a0 : float = None, + E0 : float = None, + phi0 : float = None, + zeta : float = None, + beta : float = None, + phi2 : float = None, + fill_in : bool = True, + name : str = None, + **kw): assert E0 is not None or a0 is not None, 'One of E0 or a0 must be speficied' k0 = 2.*math.pi/wavelength if E0 is None: - E0 = a0*_get_constants().m_e*_get_constants().c**2*k0/_get_constants().q_e + self.E0 = a0*_get_constants().m_e*_get_constants().c**2*k0/_get_constants().q_e if a0 is None: - a0 = E0/(_get_constants().m_e*_get_constants().c**2*k0/_get_constants().q_e) - - self.wavelength = wavelength - self.k0 = k0 - self.waist = waist - self.duration = duration - self.focal_position = focal_position - self.centroid_position = centroid_position - self.propagation_direction = propagation_direction - self.polarization_direction = polarization_direction - self.a0 = a0 - self.E0 = E0 - self.phi0 = phi0 - self.zeta = zeta - self.beta = beta - self.phi2 = phi2 - self.name = name - self.fill_in = fill_in + self.a0 = E0/(_get_constants().m_e*_get_constants().c**2*k0/_get_constants().q_e) self.handle_init(kw) +@typechecked class PICMI_AnalyticLaser(_ClassWithInit): """ Specifies a laser with an analytically described distribution @@ -117,45 +109,24 @@ class PICMI_AnalyticLaser(_ClassWithInit): Specify either amax or Emax (Emax takes precedence). - fill_in=True: Flags whether to fill in the empty spaced opened up when the grid moves """ - def __init__(self, field_expression, - wavelength, - propagation_direction = [0., 0., 1.], - polarization_direction = [1., 0., 0.], - amax = None, - Emax = None, - name = None, - fill_in = True, - **kw): + @autoargs(exclude=['kw']) + def __init__(self, field_expression : picmi_types.Expression, + wavelength : float, + propagation_direction : picmi_types.VectorFloat3 = [0., 0., 1.], + polarization_direction : picmi_types.VectorFloat3 = [1., 0., 0.], + amax : float = None, + Emax : float = None, + name : str = None, + fill_in : bool = True, + **kw): assert Emax is not None or amax is not None, 'One of Emax or amax must be speficied' k0 = 2.*math.pi/wavelength if Emax is None: - Emax = amax*_get_constants().m_e*_get_constants().c**2*k0/_get_constants().q_e + self.Emax = amax*_get_constants().m_e*_get_constants().c**2*k0/_get_constants().q_e if amax is None: - amax = Emax/(_get_constants().m_e*_get_constants().c**2*k0/_get_constants().q_e) - - self.wavelength = wavelength - self.field_expression = field_expression - self.k0 = k0 - self.propagation_direction = propagation_direction - self.polarization_direction = polarization_direction - self.amax = amax - self.Emax = Emax - self.name = name - self.fill_in = fill_in - - self.field_expression = '{}'.format(field_expression).replace('\n', '') - - # --- Find any user defined keywords in the kw dictionary. - # --- Save them and delete them from kw. - # --- It's up to the code to make sure that all parameters - # --- used in the expression are defined. - self.user_defined_kw = {} - for k in list(kw.keys()): - if re.search(r'\b%s\b'%k, self.field_expression): - self.user_defined_kw[k] = kw[k] - del kw[k] + self.amax = Emax/(_get_constants().m_e*_get_constants().c**2*k0/_get_constants().q_e) self.handle_init(kw) @@ -165,15 +136,16 @@ def __init__(self, field_expression, # ------------------ +@typechecked class PICMI_LaserAntenna(_ClassWithInit): """ Specifies the laser antenna injection method - position: Position of antenna launching the laser (vector) [m] - normal_vector: Vector normal to antenna plane (vector) [1] """ - def __init__(self, position, normal_vector, **kw): - - self.position = position - self.normal_vector = normal_vector + @autoargs(exclude=['kw']) + def __init__(self, position : picmi_types.VectorFloat3, + normal_vector : picmi_types.VectorFloat3, + **kw): self.handle_init(kw) From 0b8fa0df8321c07906f04d20465d3e01a1189bf5 Mon Sep 17 00:00:00 2001 From: Dave Grote Date: Tue, 22 Mar 2022 10:37:00 -0700 Subject: [PATCH 09/13] Added type checking to simulations, fixes in picmi_types --- PICMI_Python/picmi_types.py | 37 +++++++++++++++++++---- PICMI_Python/simulation.py | 60 +++++++++++++++++++++++-------------- 2 files changed, 68 insertions(+), 29 deletions(-) diff --git a/PICMI_Python/picmi_types.py b/PICMI_Python/picmi_types.py index e5a1936..afecc15 100644 --- a/PICMI_Python/picmi_types.py +++ b/PICMI_Python/picmi_types.py @@ -12,13 +12,38 @@ # These must be defined as strings to avoid circular definitions GridType = """typing.NewType('GridType', typing.Union[fields.PICMI_Cartesian1DGrid, - fields.PICMI_CylindricalGrid, - fields.PICMI_Cartesian2DGrid, - fields.PICMI_Cartesian3DGrid])""" + fields.PICMI_CylindricalGrid, + fields.PICMI_Cartesian2DGrid, + fields.PICMI_Cartesian3DGrid])""" -SpeciesType = """typing.NewType('SpeciesType', typing.Union[particles.PICMI_Species, - particles.PICMI_MultiSpecies])""" +_species_union = "typing.Union[particles.PICMI_Species, particles.PICMI_MultiSpecies]" -SpeciesArgument = """typing.NewType('SpeciesArgument', typing.Union[SpeciesType, Sequence[SpeciesType]])""" +SpeciesType = f"""typing.NewType('SpeciesType', {_species_union})""" +# Note that in the strings, types defined in this file cannot be used, so the _species_union needs to +# be included explicitly. +# The SpeciesArgument allows either a single species or a sequence of species. +SpeciesArgument = f"""typing.NewType('SpeciesArgument', typing.Union[{_species_union}, + Sequence[{_species_union}]])""" + +LayoutType = """typing.NewType('LayoutType', typing.Union[particles.PICMI_GriddedLayout, + particles.PICMI_PseudoRandomLayout])""" + +SolverType = """typing.NewType('SolverType', typing.Union[fields.PICMI_ElectromagneticSolver, + fields.PICMI_ElectrostaticSolver, + fields.PICMI_MagnetostaticSolver])""" + +LaserType = """typing.NewType('LaserType', typing.Union[lasers.PICMI_GaussianLaser, + lasers.PICMI_AnalyticLaser])""" + +LaserInjectionType = """typing.NewType('LaserInjectionType', lasers.PICMI_LaserAntenna)""" + +AppliedFieldType = """typing.NewType('AppliedFieldType', typing.Union[applied_fields.PICMI_ConstantAppliedField, + applied_fields.PICMI_AnalyticAppliedField])""" + +DiagnosticType = """typing.NewType('DiagnosticType', typing.Union[diagnostics.PICMI_FieldDiagnostic, + diagnostics.PICMI_ElectrostaticFieldDiagnostic, + diagnostics.PICMI_ParticleDiagnostic, + diagnostics.PICMI_LabFrameFieldDiagnostic, + diagnostics.PICMI_LabFrameParticleDiagnostic])""" diff --git a/PICMI_Python/simulation.py b/PICMI_Python/simulation.py index 445a315..23a2e1b 100644 --- a/PICMI_Python/simulation.py +++ b/PICMI_Python/simulation.py @@ -3,13 +3,24 @@ """ import math import sys +import typing + +from autoclass import autoargs +from typeguard import typechecked from .base import _ClassWithInit +from . import fields +from . import particles +from . import diagnostics +from . import lasers +from . import diagnostics +from . import picmi_types # --------------------- # Main simulation object # --------------------- +@typechecked class PICMI_Simulation(_ClassWithInit): """ Creates a Simulation object @@ -47,16 +58,17 @@ class PICMI_Simulation(_ClassWithInit): Code specific arguments ; should be prefixed with the `codename` """ - def __init__(self, solver=None, time_step_size=None, max_steps=None, max_time=None, verbose=None, - particle_shape='linear', gamma_boost=None, cpu_split=None, load_balancing=None, **kw): - - self.solver = solver - self.time_step_size = time_step_size - self.verbose = verbose - self.max_steps = max_steps - self.max_time = max_time - self.particle_shape = particle_shape - self.gamma_boost = gamma_boost + @autoargs(exclude=['kw']) + def __init__(self, solver : picmi_types.SolverType = None, + time_step_size : float = None, + max_steps : int = None, + max_time : float = None, + verbose : int = None, + particle_shape : str = 'linear', + gamma_boost : float = None, + cpu_split : bool = None, + load_balancing : bool = None, + **kw): self.species = [] self.layouts = [] @@ -71,12 +83,11 @@ def __init__(self, solver=None, time_step_size=None, max_steps=None, max_time=No self.diagnostics = [] - self.cpu_split = cpu_split - self.load_balancing = load_balancing - self.handle_init(kw) - def add_species(self, species, layout, initialize_self_field=None): + def add_species(self, species : picmi_types.SpeciesType, + layout : picmi_types.LayoutType, + initialize_self_field : bool = None): """ Add species to be used in the simulation @@ -102,9 +113,11 @@ def add_species(self, species, layout, initialize_self_field=None): self.injection_plane_normal_vectors.append(None) - def add_species_through_plane(self, species, layout, - injection_plane_position, injection_plane_normal_vector, - initialize_self_field=None ): + def add_species_through_plane(self, species : picmi_types.SpeciesType, + layout : picmi_types.LayoutType, + injection_plane_position : picmi_types.VectorFloat3, + injection_plane_normal_vector : picmi_types.VectorFloat3, + initialize_self_field : bool = None ): """ Add species to be used in the simulation @@ -122,7 +135,8 @@ def add_species_through_plane(self, species, layout, self.injection_plane_normal_vectors.append(injection_plane_normal_vector) - def add_laser(self, laser, injection_method): + def add_laser(self, laser : picmi_types.LaserType, + injection_method : picmi_types.LaserInjectionType): """ Add a laser pulses that to be injected in the simulation @@ -144,7 +158,7 @@ def add_laser(self, laser, injection_method): self.lasers.append(laser) self.laser_injection_methods.append(injection_method) - def add_applied_field(self, applied_field): + def add_applied_field(self, applied_field : picmi_types.AppliedFieldType): """ Add an applied field @@ -156,7 +170,7 @@ def add_applied_field(self, applied_field): """ self.applied_fields.append(applied_field) - def add_diagnostic(self, diagnostic): + def add_diagnostic(self, diagnostic : picmi_types.DiagnosticType): """ Add a diagnostic - diagnostic: object @@ -164,7 +178,7 @@ def add_diagnostic(self, diagnostic): """ self.diagnostics.append(diagnostic) - def set_max_step(self, max_steps): + def set_max_step(self, max_steps : int): """ Set the default number of steps for the simulation (i.e. the number of steps that gets written when calling `write_input_file`) @@ -179,7 +193,7 @@ def set_max_step(self, max_steps): """ self.max_steps = max_steps - def write_input_file(self, file_name): + def write_input_file(self, file_name : str): """ Write the parameters of the simulation, as defined in the PICMI input, into another, more code-specific input file. @@ -194,7 +208,7 @@ def write_input_file(self, file_name): """ raise NotImplementedError - def step(self, nsteps=1): + def step(self, nsteps : int = 1): """ Run the simulation for `nsteps` timesteps From 9fff6503fda8b78535ad14a7e6ca5c8da665b3e2 Mon Sep 17 00:00:00 2001 From: Dave Grote Date: Wed, 23 Mar 2022 15:51:53 -0700 Subject: [PATCH 10/13] Added typing to particles.py --- PICMI_Python/base.py | 25 +++- PICMI_Python/lasers.py | 1 + PICMI_Python/particles.py | 261 ++++++++++++++++-------------------- PICMI_Python/picmi_types.py | 10 ++ 4 files changed, 147 insertions(+), 150 deletions(-) diff --git a/PICMI_Python/base.py b/PICMI_Python/base.py index 31346b1..76bbefb 100644 --- a/PICMI_Python/base.py +++ b/PICMI_Python/base.py @@ -38,10 +38,18 @@ def _get_constants(): class _ClassWithInit(object): def _check_vector_lengths(self): for arg_name, arg_type in self.__init__.__annotations__.items(): - if arg_type in [picmi_types.VectorFloat3, picmi_types.VectorInt3]: + if arg_type in [picmi_types.VectorFloat3, picmi_types.VectorInt3, picmi_types.VectorExpression3]: arg_value = getattr(self, arg_name) assert len(arg_value) == 3, Exception(f'{arg_name} must have a length of 3') + def _add_to_user_defined_kw(self, arg_value, kw): + # --- The dictionary is created if needed + self.user_defined_kw = getattr(self, 'user_defined_kw', {}) + if arg_value is not None: + for k in list(kw.keys()): + if re.search(r'\b%s\b'%k, arg_value): + self.user_defined_kw[k] = kw.pop(k) + def _process_expression_arguments(self, kw): """For arguments that are of type Expression, save any keyword arguments used in the expression in the user_defined_kw dictionary. @@ -56,12 +64,17 @@ def _process_expression_arguments(self, kw): arg_value = arg_value.replace('\n', '') setattr(self, arg_name, arg_value) - # --- The dictionary is created if needed - self.user_defined_kw = getattr(self, 'user_defined_kw', {}) + self._add_to_user_defined_kw(arg_value, kw) + + elif arg_type == picmi_types.VectorExpression3: + + arg_value = getattr(self, arg_name) + + # --- Remove any line feeds from the expressions if arg_value is not None: - for k in list(kw.keys()): - if re.search(r'\b%s\b'%k, arg_value): - self.user_defined_kw[k] = kw.pop(k) + for i in range(3): + arg_value[i] = arg_value[i].replace('\n', '') + self._add_to_user_defined_kw(arg_value[i], kw) def handle_init(self, kw): self._check_vector_lengths() diff --git a/PICMI_Python/lasers.py b/PICMI_Python/lasers.py index bf7fbd6..1a3e01c 100644 --- a/PICMI_Python/lasers.py +++ b/PICMI_Python/lasers.py @@ -4,6 +4,7 @@ import math import sys import re +import typing from autoclass import autoargs from typeguard import typechecked diff --git a/PICMI_Python/particles.py b/PICMI_Python/particles.py index e0f1bbe..ce3f1c7 100644 --- a/PICMI_Python/particles.py +++ b/PICMI_Python/particles.py @@ -5,15 +5,24 @@ import math import sys import re +import typing +from collections.abc import Sequence + +from autoclass import autoargs +from typeguard import typechecked + import numpy as np from .base import _ClassWithInit +from . import fields +from . import picmi_types # --------------- # Physics objects # --------------- +@typechecked class PICMI_Species(_ClassWithInit): """ Species @@ -31,23 +40,20 @@ class PICMI_Species(_ClassWithInit): methods_list = ['Boris' , 'Vay', 'Higuera-Cary', 'Li', 'free-streaming', 'LLRK4'] - def __init__(self, particle_type=None, name=None, charge_state=None, charge=None, mass=None, - initial_distribution=None, particle_shape=None, density_scale=None, method=None, **kw): - + @autoargs(exclude=['kw']) + def __init__(self, particle_type : str = None, + name : str = None, + charge_state : float = None, + charge : float = None, + mass : float = None, + initial_distribution : picmi_types.DistributionType = None, + particle_shape : str = None, + density_scale : float = None, + method : str = None, + **kw): assert method is None or method in PICMI_Species.methods_list or method.startswith('other:'), \ Exception('method must starts with either "other:", or be one of the following '+', '.join(PICMI_Species.methods_list)) - - self.method = method - self.particle_type = particle_type - self.name = name - self.charge = charge - self.charge_state = charge_state - self.mass = mass - self.initial_distribution = initial_distribution - self.particle_shape = particle_shape - self.density_scale = density_scale - self.interactions = [] self.handle_init(kw) @@ -56,9 +62,11 @@ def activate_field_ionization(self, model, product_species): # --- TODO: One way of handling interactions is to add a class for each type # --- of interaction. Instances would be added to the interactions list # --- instead of the list of parameters. + # --- This interface is not yet defined. self.interactions.append(['ionization', model, product_species]) +@typechecked class PICMI_MultiSpecies(_ClassWithInit): """ INCOMPLETE: proportions argument is not implemented @@ -78,37 +86,35 @@ class PICMI_MultiSpecies(_ClassWithInit): # --- defined in the codes PICMI implementation. Species_class = None - def __init__(self, particle_types=None, names=None, charge_states=None, charges=None, masses=None, - proportions=None, initial_distribution=None, particle_shape=None, - **kw): - - self.particle_types = particle_types - self.names = names - self.charges = charges - self.charge_states = charge_states - self.masses = masses - self.proportions = proportions - self.initial_distribution = initial_distribution - self.particle_shape = particle_shape + @autoargs(exclude=['kw']) + def __init__(self, particle_types : str = None, + names : str = None, + charge_states : Sequence[int] = None, + charges : Sequence[float] = None, + masses : Sequence[float] = None, + proportions : Sequence[float] = None, + initial_distribution : picmi_types.DistributionType = None, + particle_shape : str = None, + **kw): self.nspecies = None - self.check_nspecies(particle_types) - self.check_nspecies(names) - self.check_nspecies(charges) - self.check_nspecies(charge_states) - self.check_nspecies(masses) - self.check_nspecies(proportions) + self._check_nspecies(particle_types) + self._check_nspecies(names) + self._check_nspecies(charges) + self._check_nspecies(charge_states) + self._check_nspecies(masses) + self._check_nspecies(proportions) # --- Create the instances of each species self.species_instances_list = [] self.species_instances_dict = {} for i in range(self.nspecies): - particle_type = self.get_input_item(particle_types, i) - name = self.get_input_item(names, i) - charge = self.get_input_item(charges, i) - charge_state = self.get_input_item(charge_states, i) - mass = self.get_input_item(masses, i) - proportion = self.get_input_item(proportions, i) + particle_type = self._get_input_item(particle_types, i) + name = self._get_input_item(names, i) + charge = self._get_input_item(charges, i) + charge_state = self._get_input_item(charge_states, i) + mass = self._get_input_item(masses, i) + proportion = self._get_input_item(proportions, i) specie = PICMI_MultiSpecies.Species_class(particle_type = particle_type, name = name, charge = charge, @@ -122,7 +128,7 @@ def __init__(self, particle_types=None, names=None, charge_states=None, charges= self.handle_init(kw) - def check_nspecies(self, var): + def _check_nspecies(self, var): if var is not None: try: nvars = len(var) @@ -131,7 +137,7 @@ def check_nspecies(self, var): assert self.nspecies is None or self.nspecies == nvars, Exception('All inputs must have the same length') self.nspecies = nvars - def get_input_item(self, var, i): + def _get_input_item(self, var, i): if var is None: return None else: @@ -152,6 +158,7 @@ def __getitem__(self, key): return self.species_instances_list[key] +@typechecked class PICMI_GaussianBunchDistribution(_ClassWithInit): """ Describes a Gaussian distribution of particles @@ -162,22 +169,19 @@ class PICMI_GaussianBunchDistribution(_ClassWithInit): - centroid_velocity=[0,0,0]: Velocity (gamma*V) of the bunch centroid at t=0 (vector) [m/s] - velocity_divergence=[0,0,0]: Expansion rate of the bunch at t=0 (vector) [m/s/m] """ - def __init__(self,n_physical_particles, rms_bunch_size, - rms_velocity = [0.,0.,0.], - centroid_position = [0.,0.,0.], - centroid_velocity = [0.,0.,0.], - velocity_divergence = [0.,0.,0.], - **kw): - self.n_physical_particles = n_physical_particles - self.rms_bunch_size = rms_bunch_size - self.rms_velocity = rms_velocity - self.centroid_position = centroid_position - self.centroid_velocity = centroid_velocity - self.velocity_divergence = velocity_divergence + @autoargs(exclude=['kw']) + def __init__(self,n_physical_particles : int, + rms_bunch_size : picmi_types.VectorFloat3, + rms_velocity : picmi_types.VectorFloat3 = [0.,0.,0.], + centroid_position : picmi_types.VectorFloat3 = [0.,0.,0.], + centroid_velocity : picmi_types.VectorFloat3 = [0.,0.,0.], + velocity_divergence : picmi_types.VectorFloat3 = [0.,0.,0.], + **kw): self.handle_init(kw) +@typechecked class PICMI_UniformDistribution(_ClassWithInit): """ Describes a uniform density distribution of particles @@ -189,23 +193,19 @@ class PICMI_UniformDistribution(_ClassWithInit): - fill_in: Flags whether to fill in the empty spaced opened up when the grid moves """ - def __init__(self, density, - lower_bound = [None,None,None], - upper_bound = [None,None,None], - rms_velocity = [0.,0.,0.], - directed_velocity = [0.,0.,0.], - fill_in = None, - **kw): - self.density = density - self.lower_bound = lower_bound - self.upper_bound = upper_bound - self.rms_velocity = rms_velocity - self.directed_velocity = directed_velocity - self.fill_in = fill_in + @autoargs(exclude=['kw']) + def __init__(self, density : float, + lower_bound : picmi_types.VectorFloatNone3 = [None,None,None], + upper_bound : picmi_types.VectorFloatNone3 = [None,None,None], + rms_velocity : picmi_types.VectorFloat3 = [0.,0.,0.], + directed_velocity : picmi_types.VectorFloat3 = [0.,0.,0.], + fill_in : bool = None, + **kw): self.handle_init(kw) +@typechecked class PICMI_AnalyticDistribution(_ClassWithInit): """ Describes a uniform density plasma @@ -229,52 +229,20 @@ class PICMI_AnalyticDistribution(_ClassWithInit): ...) """ - def __init__(self, density_expression, - momentum_expressions = [None, None, None], - lower_bound = [None,None,None], - upper_bound = [None,None,None], - rms_velocity = [0.,0.,0.], - directed_velocity = [0.,0.,0.], - fill_in = None, - **kw): - self.density_expression = '{}'.format(density_expression).replace('\n', '') - self.momentum_expressions = momentum_expressions - self.lower_bound = lower_bound - self.upper_bound = upper_bound - self.rms_velocity = rms_velocity - self.directed_velocity = directed_velocity - self.fill_in = fill_in - - # --- Convert momentum_expressions to string if needed. - if self.momentum_expressions[0] is not None: - self.momentum_expressions[0] = '{}'.format(self.momentum_expressions[0]).replace('\n', '') - if self.momentum_expressions[1] is not None: - self.momentum_expressions[1] = '{}'.format(self.momentum_expressions[1]).replace('\n', '') - if self.momentum_expressions[2] is not None: - self.momentum_expressions[2] = '{}'.format(self.momentum_expressions[2]).replace('\n', '') - - # --- Find any user defined keywords in the kw dictionary. - # --- Save them and delete them from kw. - # --- It's up to the code to make sure that all parameters - # --- used in the expression are defined. - self.user_defined_kw = {} - for k in list(kw.keys()): - if re.search(r'\b%s\b'%k, self.density_expression): - self.user_defined_kw[k] = kw[k] - del kw[k] - elif self.momentum_expressions[0] is not None and re.search(r'\b%s\b'%k, self.momentum_expressions[0]): - self.user_defined_kw[k] = kw[k] - del kw[k] - elif self.momentum_expressions[1] is not None and re.search(r'\b%s\b'%k, self.momentum_expressions[1]): - self.user_defined_kw[k] = kw[k] - del kw[k] - elif self.momentum_expressions[2] is not None and re.search(r'\b%s\b'%k, self.momentum_expressions[2]): - self.user_defined_kw[k] = kw[k] - del kw[k] + @autoargs(exclude=['kw']) + def __init__(self, density_expression : picmi_types.Expression, + momentum_expressions : picmi_types.VectorExpression3 = [None, None, None], + lower_bound : picmi_types.VectorFloatNone3 = [None,None,None], + upper_bound : picmi_types.VectorFloatNone3 = [None,None,None], + rms_velocity : picmi_types.VectorFloat3 = [0.,0.,0.], + directed_velocity : picmi_types.VectorFloat3 = [0.,0.,0.], + fill_in : bool = None, + **kw): self.handle_init(kw) +@typechecked class PICMI_ParticleListDistribution(_ClassWithInit): """ Load particles at the specified positions and velocities @@ -286,8 +254,16 @@ class PICMI_ParticleListDistribution(_ClassWithInit): - uz=0.: List of uz positions of the particles (uz = gamma*vz) [m/s] - weight: Particle weight or list of weights, number of real particles per simulation particle """ - def __init__(self, x=0., y=0., z=0., ux=0., uy=0., uz=0., weight=0., - **kw): + @autoargs(exclude=['kw']) + def __init__(self, x : Sequence[float] = 0., + y : Sequence[float] = 0., + z : Sequence[float] = 0., + ux : Sequence[float] = 0., + uy : Sequence[float] = 0., + uz : Sequence[float] = 0., + weight : typing.Union[float, Sequence[float]] = 0., + **kw): + # --- Get length of arrays, set to one for scalars lenx = np.size(x) leny = np.size(y) @@ -298,36 +274,28 @@ def __init__(self, x=0., y=0., z=0., ux=0., uy=0., uz=0., weight=0., lenw = np.size(weight) maxlen = max(lenx, leny, lenz, lenux, lenuy, lenuz, lenw) - assert lenx==maxlen or lenx==1, "Length of x doesn't match len of others" - assert leny==maxlen or leny==1, "Length of y doesn't match len of others" - assert lenz==maxlen or lenz==1, "Length of z doesn't match len of others" - assert lenux==maxlen or lenux==1, "Length of ux doesn't match len of others" - assert lenuy==maxlen or lenuy==1, "Length of uy doesn't match len of others" - assert lenuz==maxlen or lenuz==1, "Length of uz doesn't match len of others" - assert lenw==maxlen or lenw==1, "Length of weight doesn't match len of others" + assert lenx==maxlen or lenx==1, "Length of x doesn't match length of others" + assert leny==maxlen or leny==1, "Length of y doesn't match length of others" + assert lenz==maxlen or lenz==1, "Length of z doesn't match length of others" + assert lenux==maxlen or lenux==1, "Length of ux doesn't match length of others" + assert lenuy==maxlen or lenuy==1, "Length of uy doesn't match length of others" + assert lenuz==maxlen or lenuz==1, "Length of uz doesn't match length of others" + assert lenw==maxlen or lenw==1, "Length of weight doesn't match length of others" if lenx == 1: - x = np.array(x)*np.ones(maxlen) + self.x = np.array(x)*np.ones(maxlen) if leny == 1: - y = np.array(y)*np.ones(maxlen) + self.y = np.array(y)*np.ones(maxlen) if lenz == 1: - z = np.array(z)*np.ones(maxlen) + self.z = np.array(z)*np.ones(maxlen) if lenux == 1: - ux = np.array(ux)*np.ones(maxlen) + self.ux = np.array(ux)*np.ones(maxlen) if lenuy == 1: - uy = np.array(uy)*np.ones(maxlen) + self.uy = np.array(uy)*np.ones(maxlen) if lenuz == 1: - uz = np.array(uz)*np.ones(maxlen,'d') + self.uz = np.array(uz)*np.ones(maxlen,'d') # --- Note that weight can be a scalar - self.weight = weight - self.x = x - self.y = y - self.z = z - self.ux = ux - self.uy = uy - self.uz = uz - self.handle_init(kw) @@ -336,6 +304,7 @@ def __init__(self, x=0., y=0., z=0., ux=0., uy=0., uz=0., weight=0., # ------------------ +@typechecked class PICMI_ParticleDistributionPlanarInjector(_ClassWithInit): """ Describes the injection of particles from a plane @@ -344,15 +313,17 @@ class PICMI_ParticleDistributionPlanarInjector(_ClassWithInit): - plane_velocity: Velocity of the plane of injection (vector) [m/s] - method: InPlace - method of injection. One of 'InPlace', or 'Plane' """ - def __init__(self, position, plane_normal, plane_velocity=[0.,0.,0.], method='InPlace', **kw): - self.position = position - self.plane_normal = plane_normal - self.plane_velocity = plane_velocity - self.method = method + @autoargs(exclude=['kw']) + def __init__(self, position : picmi_types.VectorFloat3, + plane_normal : picmi_types.VectorFloat3, + plane_velocity : picmi_types.VectorFloat3 = [0.,0.,0.], + method : str = 'InPlace', + **kw): self.handle_init(kw) +@typechecked class PICMI_GriddedLayout(_ClassWithInit): """ Specifies a gridded layout of particles @@ -360,13 +331,15 @@ class PICMI_GriddedLayout(_ClassWithInit): - grid: grid object specifying the grid to follow (optional) If not specified, the underlying grid of the code is used. """ - def __init__(self, n_macroparticle_per_cell, grid=None, **kw): - self.n_macroparticle_per_cell = n_macroparticle_per_cell - self.grid = grid + @autoargs(exclude=['kw']) + def __init__(self, n_macroparticle_per_cell : picmi_types.VectorInt3, + grid : picmi_types.GridType = None, + **kw): self.handle_init(kw) +@typechecked class PICMI_PseudoRandomLayout(_ClassWithInit): """ Specifies a pseudo-random layout of the particles @@ -378,14 +351,14 @@ class PICMI_PseudoRandomLayout(_ClassWithInit): - grid: grid object specifying the grid to follow for n_macroparticles_per_cell (optional) If not specified, the underlying grid of the code is used. """ - def __init__(self, n_macroparticles=None, n_macroparticles_per_cell=None, seed=None, grid=None, **kw): + @autoargs(exclude=['kw']) + def __init__(self, n_macroparticles : int = None, + n_macroparticles_per_cell : int = None, + seed : int = None, + grid : picmi_types.GridType = None, + **kw): assert (n_macroparticles is not None)^(n_macroparticles_per_cell is not None), \ Exception('Only one of n_macroparticles and n_macroparticles_per_cell must be specified') - self.n_macroparticles = n_macroparticles - self.n_macroparticles_per_cell = n_macroparticles_per_cell - self.seed = seed - self.grid = grid - self.handle_init(kw) diff --git a/PICMI_Python/picmi_types.py b/PICMI_Python/picmi_types.py index afecc15..2e56d3d 100644 --- a/PICMI_Python/picmi_types.py +++ b/PICMI_Python/picmi_types.py @@ -6,8 +6,10 @@ import numpy as np VectorFloat3 = typing.NewType('VectorFloat3', typing.Union[Sequence[float], np.ndarray[3, np.float64]]) +VectorFloatNone3 = typing.NewType('VectorFloat3', typing.Union[Sequence[typing.Union[float, None]], np.ndarray[3, np.float64]]) VectorInt3 = typing.NewType('VectorInt3', typing.Union[Sequence[int], np.ndarray[3, np.int64]]) Expression = typing.NewType('Expression', str) +VectorExpression3 = typing.NewType('VectorExpression3', Sequence[str]) # These must be defined as strings to avoid circular definitions @@ -26,6 +28,14 @@ SpeciesArgument = f"""typing.NewType('SpeciesArgument', typing.Union[{_species_union}, Sequence[{_species_union}]])""" +# Note that this is only used in particles.py and uses classes defined there. +# The "particles" prefix can not be used on the classes since the module +# should not reference itself. +DistributionType = """typing.NewType('DistributionType', typing.Union[PICMI_GaussianBunchDistribution, + PICMI_UniformDistribution, + PICMI_AnalyticDistribution, + PICMI_ParticleListDistribution])""" + LayoutType = """typing.NewType('LayoutType', typing.Union[particles.PICMI_GriddedLayout, particles.PICMI_PseudoRandomLayout])""" From 8c33818fe76cd37c433c6b60d42fe38701b64eff Mon Sep 17 00:00:00 2001 From: Dave Grote Date: Thu, 24 Mar 2022 19:15:11 -0700 Subject: [PATCH 11/13] Added types to fields.py --- PICMI_Python/base.py | 3 +- PICMI_Python/fields.py | 579 +++++++++++++++++++----------------- PICMI_Python/picmi_types.py | 19 +- 3 files changed, 320 insertions(+), 281 deletions(-) diff --git a/PICMI_Python/base.py b/PICMI_Python/base.py index 76bbefb..a2094bd 100644 --- a/PICMI_Python/base.py +++ b/PICMI_Python/base.py @@ -40,7 +40,8 @@ def _check_vector_lengths(self): for arg_name, arg_type in self.__init__.__annotations__.items(): if arg_type in [picmi_types.VectorFloat3, picmi_types.VectorInt3, picmi_types.VectorExpression3]: arg_value = getattr(self, arg_name) - assert len(arg_value) == 3, Exception(f'{arg_name} must have a length of 3') + if arg_value is not None: + assert len(arg_value) == 3, Exception(f'{arg_name} must have a length of 3') def _add_to_user_defined_kw(self, arg_value, kw): # --- The dictionary is created if needed diff --git a/PICMI_Python/fields.py b/PICMI_Python/fields.py index 31672ca..f41ddb4 100644 --- a/PICMI_Python/fields.py +++ b/PICMI_Python/fields.py @@ -3,13 +3,20 @@ """ import math import sys +import typing +from collections.abc import Sequence + +from autoclass import autoargs +from typeguard import typechecked from .base import _ClassWithInit +from . import picmi_types # --------------- # Physics objects # --------------- +@typechecked class PICMI_ElectromagneticSolver(_ClassWithInit): """ Electromagnetic field solver @@ -30,31 +37,29 @@ class PICMI_ElectromagneticSolver(_ClassWithInit): methods_list = ['Yee', 'CKC', 'Lehe', 'PSTD', 'PSATD', 'GPSTD', 'ECT'] - def __init__(self, grid, method=None, stencil_order=None, cfl=None, l_nodal=None, - source_smoother=None, field_smoother=None, subcycling=None, - galilean_velocity=None, divE_cleaning=None, divB_cleaning=None, - pml_divE_cleaning=None, pml_divB_cleaning=None, **kw): + @autoargs(exclude=['kw']) + def __init__(self, grid : picmi_types.FieldsGridType, + method : str = None, + stencil_order : picmi_types.VectorInt3 = None, + cfl : float = None, + l_nodal : bool = None, + source_smoother : picmi_types.SmootherType = None, + field_smoother : picmi_types.SmootherType = None, + subcycling : int = None, + galilean_velocity : picmi_types.VectorFloat3 = None, + divE_cleaning : bool = None, + divB_cleaning : bool = None, + pml_divE_cleaning : bool = None, + pml_divB_cleaning : bool = None, + **kw): assert method is None or method in PICMI_ElectromagneticSolver.methods_list, \ Exception('method must be one of '+', '.join(PICMI_ElectromagneticSolver.methods_list)) - self.grid = grid - self.method = method - self.cfl = cfl - self.stencil_order = stencil_order - self.l_nodal = l_nodal - self.source_smoother = source_smoother - self.field_smoother = field_smoother - self.subcycling = subcycling - self.galilean_velocity = galilean_velocity - self.divE_cleaning = divE_cleaning - self.divB_cleaning = divB_cleaning - self.pml_divE_cleaning = pml_divE_cleaning - self.pml_divB_cleaning = pml_divB_cleaning - self.handle_init(kw) +@typechecked class PICMI_ElectrostaticSolver(_ClassWithInit): """ Electrostatic field solver @@ -66,20 +71,20 @@ class PICMI_ElectrostaticSolver(_ClassWithInit): methods_list = ['FFT', 'Multigrid'] - def __init__(self, grid, method=None, - required_precision=None, maximum_iterations=None, **kw): + @autoargs(exclude=['kw']) + def __init__(self, grid : picmi_types.FieldsGridType, + method : str = None, + required_precision : float = None, + maximum_iterations : int = None, + **kw): assert method is None or method in PICMI_ElectrostaticSolver.methods_list, \ Exception('method must be one of '+', '.join(PICMI_ElectrostaticSolver.methods_list)) - self.grid = grid - self.method = method - self.required_precision = required_precision - self.maximum_iterations = maximum_iterations - self.handle_init(kw) +@typechecked class PICMI_MagnetostaticSolver(_ClassWithInit): """ Magnetostatic field solver @@ -89,14 +94,14 @@ class PICMI_MagnetostaticSolver(_ClassWithInit): methods_list = ['FFT', 'Multigrid'] - def __init__(self, grid, method=None, **kw): + @autoargs(exclude=['kw']) + def __init__(self, grid : picmi_types.FieldsGridType, + method : str = None, + **kw): assert method is None or method in PICMI_MagnetostaticSolver.methods_list, \ Exception('method must be one of '+', '.join(PICMI_MagnetostaticSolver.methods_list)) - self.grid = grid - self.method = method - self.handle_init(kw) @@ -105,6 +110,7 @@ def __init__(self, grid, method=None, **kw): # ------------------ +@typechecked class PICMI_BinomialSmoother(_ClassWithInit): """ Descibes a binomial smoother operator (applied to grids) @@ -113,15 +119,17 @@ class PICMI_BinomialSmoother(_ClassWithInit): - stride: Stride along each axis. (vector) - alpha: Smoothing coefficients along each axis. (vector) """ - def __init__(self, n_pass=None, compensation=None, stride=None, alpha=None, **kw): - self.n_pass = n_pass - self.compensation = compensation - self.stride = stride - self.alpha = alpha + @autoargs(exclude=['kw']) + def __init__(self, n_pass : picmi_types.VectorInt3 = None, + compensation : picmi_types.VectorBool3 = None, + stride : picmi_types.VectorInt3 = None, + alpha : picmi_types.VectorFloat3 = None, + **kw): self.handle_init(kw) +@typechecked class PICMI_Cartesian1DGrid(_ClassWithInit): """ One-dimensional Cartesian grid @@ -144,8 +152,8 @@ class PICMI_Cartesian1DGrid(_ClassWithInit): - refined_regions: List of refined regions, each element being a list of the format [level, lo, hi, refinement_factor], with level being the refinement level, with 1 being the first level of refinement, 2 being the second etc, - lo and hi being vectors of length 2 specifying the extent of the region, - and refinement_factor defaulting to [2,2] (relative to next lower level) + lo and hi being vectors of length 1 specifying the extent of the region, + and refinement_factor defaulting to [2] (relative to next lower level) - lower_bound_particles: Position of particle lower bound (vector of floats) [m] - upper_bound_particles: Position of particle upper bound (vector of floats) [m] @@ -165,15 +173,30 @@ class PICMI_Cartesian1DGrid(_ClassWithInit): number_of_dimensions = 1 - def __init__(self, number_of_cells=None, lower_bound=None, upper_bound=None, - lower_boundary_conditions=None, upper_boundary_conditions=None, - nx=None, xmin=None, xmax=None, bc_xmin=None, bc_xmax=None, - moving_window_velocity=None, refined_regions=[],lower_bound_particles=None, upper_bound_particles=None, - xmin_particles=None, xmax_particles=None, - lower_boundary_conditions_particles=None, upper_boundary_conditions_particles=None, - bc_xmin_particles=None, bc_xmax_particles=None, - guard_cells=None, pml_cells=None, - **kw): + @autoargs(exclude=['kw']) + def __init__(self, number_of_cells : picmi_types.VectorInt1 = None, + lower_bound : picmi_types.VectorFloat1 = None, + upper_bound : picmi_types.VectorFloat1 = None, + lower_boundary_conditions : picmi_types.VectorString1 = None, + upper_boundary_conditions : picmi_types.VectorString1 = None, + nx : int = None, + xmin : float = None, + xmax : float = None, + bc_xmin : str = None, + bc_xmax : str = None, + moving_window_velocity : picmi_types.VectorFloat1 = None, + refined_regions : Sequence[list] = None, + lower_bound_particles : picmi_types.VectorFloat1 = None, + upper_bound_particles : picmi_types.VectorFloat1 = None, + xmin_particles : float = None, + xmax_particles : float = None, + lower_boundary_conditions_particles : picmi_types.VectorString1 = None, + upper_boundary_conditions_particles : picmi_types.VectorString1 = None, + bc_xmin_particles : str = None, + bc_xmax_particles : str = None, + guard_cells : picmi_types.VectorInt1 = None, + pml_cells : picmi_types.VectorInt1 = None, + **kw): # Sanity check and init of input arguments related to grid parameters assert (number_of_cells is None) and (nx is not None) or \ @@ -193,86 +216,66 @@ def __init__(self, number_of_cells=None, lower_bound=None, upper_bound=None, Exception('Either upper_boundary_conditions or bc_xmax must be specified') if number_of_cells is None: - number_of_cells = [nx] + self.number_of_cells = [nx] if lower_bound is None: - lower_bound = [xmin] + self.lower_bound = [xmin] if upper_bound is None: - upper_bound = [xmax] + self.upper_bound = [xmax] if lower_boundary_conditions is None: - lower_boundary_conditions = [bc_xmin] + self.lower_boundary_conditions = [bc_xmin] if upper_boundary_conditions is None: - upper_boundary_conditions = [bc_xmax,] + self.upper_boundary_conditions = [bc_xmax,] # Sanity check and init of input arguments related to particle boundary parameters # By default, if not specified, particle boundary values are the same as field boundary values # By default, if not specified, particle boundary conditions are the same as field boundary conditions if lower_bound_particles is None: if (xmin_particles is None): - lower_bound_particles = lower_bound + self.lower_bound_particles = lower_bound else: - lower_bound_particles = [xmin_particles] + self.lower_bound_particles = [xmin_particles] if upper_bound_particles is None: if (xmax_particles is None): - upper_bound_particles = upper_bound + self.upper_bound_particles = upper_bound else: - upper_bound_particles=[xmax_particles] + self.upper_bound_particles = [xmax_particles] if lower_boundary_conditions_particles is None: if (bc_xmin_particles is None): - lower_boundary_conditions_particles = lower_boundary_conditions + self.lower_boundary_conditions_particles = lower_boundary_conditions else: - lower_boundary_conditions_particles = [bc_xmin_particles] + self.lower_boundary_conditions_particles = [bc_xmin_particles] if upper_boundary_conditions_particles is None: if (bc_xmax_particles is None): - upper_boundary_conditions_particles = upper_boundary_conditions + self.upper_boundary_conditions_particles = upper_boundary_conditions else: - upper_boundary_conditions_particles = [bc_xmax_particles] - - # Sanity check on dimensionality of vector quantities - assert len(number_of_cells) == 1, Exception('Wrong number of cells specified') - assert len(lower_bound) == 1, Exception('Wrong number of lower bounds specified') - assert len(upper_bound) == 1, Exception('Wrong number of upper bounds specified') - assert len(lower_boundary_conditions) == 1, Exception('Wrong number of lower boundary conditions specified') - assert len(upper_boundary_conditions) == 1, Exception('Wrong number of upper boundary conditions specified') - assert len(lower_bound_particles) == 1, Exception('Wrong number of particle lower bounds specified') - assert len(upper_bound_particles) == 1, Exception('Wrong number of particle upper bounds specified') - assert len(lower_boundary_conditions_particles) == 1, Exception('Wrong number of lower particle boundary conditions specified') - assert len(upper_boundary_conditions_particles) == 1, Exception('Wrong number of upper particle boundary conditions specified') - - self.number_of_cells = number_of_cells - self.lower_bound = lower_bound - self.upper_bound = upper_bound - self.lower_boundary_conditions = lower_boundary_conditions - self.upper_boundary_conditions = upper_boundary_conditions - self.lower_bound_particles = lower_bound_particles - self.upper_bound_particles = upper_bound_particles - self.lower_boundary_conditions_particles = lower_boundary_conditions_particles - self.upper_boundary_conditions_particles = upper_boundary_conditions_particles - self.guard_cells = guard_cells - self.pml_cells = pml_cells - - self.moving_window_velocity = moving_window_velocity - - self.refined_regions = refined_regions - - for region in self.refined_regions: - if len(region) == 3: - region.append([2]) - assert len(region[1]) == 1, Exception('The lo extent of the refined region must be a vector of length 2') - assert len(region[2]) == 1, Exception('The hi extent of the refined region must be a vector of length 2') - assert len(region[3]) == 1, Exception('The refinement factor of the refined region must be a vector of length 2') + self.upper_boundary_conditions_particles = [bc_xmax_particles] + + self.refined_regions = [] + if refined_regions is not None: + for region in refined_regions: + self.add_refined_region(*region) self.handle_init(kw) - def add_refined_region(self, level, lo, hi, refinement_factor=[2]): + def add_refined_region(self, level : int, + lo : picmi_types.VectorInt1, + hi : picmi_types.VectorInt1, + refinement_factor : picmi_types.VectorInt1 = None): """Add a refined region. - level: the refinement level, with 1 being the first level of refinement, 2 being the second etc. - lo, hi: vectors of length 2 specifying the extent of the region - refinement_factor: defaulting to [2,2] (relative to next lower level) + - level: the refinement level, with 1 being the first level of refinement, 2 being the second etc. + - lo, hi: vectors of length 1 specifying the extent of the region + - refinement_factor: defaulting to [2] (relative to next lower level) """ + if refinement_factor is None: + refinement_factor = [2] + assert len(lo) == 1, Exception('The lo extent of the refined region must be a vector of length 1') + assert len(hi) == 1, Exception('The hi extent of the refined region must be a vector of length 1') + assert len(refinement_factor) == 1, Exception('The refinement factor of the refined region must be a vector of length 1') self.refined_regions.append([level, lo, hi, refinement_factor]) +@typechecked class PICMI_CylindricalGrid(_ClassWithInit): """ Axisymmetric, cylindrical grid @@ -326,17 +329,40 @@ class PICMI_CylindricalGrid(_ClassWithInit): number_of_dimensions = 2 - def __init__(self, number_of_cells=None, lower_bound=None, upper_bound=None, - lower_boundary_conditions=None, upper_boundary_conditions=None, - nr=None, nz=None, n_azimuthal_modes=None, - rmin=None, rmax=None, zmin=None, zmax=None, - bc_rmin=None, bc_rmax=None, bc_zmin=None, bc_zmax=None, - moving_window_velocity=None, refined_regions=[], - lower_bound_particles=None, upper_bound_particles=None, - rmin_particles=None, rmax_particles=None, zmin_particles=None, zmax_particles=None, - lower_boundary_conditions_particles=None, upper_boundary_conditions_particles=None, - bc_rmin_particles=None, bc_rmax_particles=None, bc_zmin_particles=None, bc_zmax_particles=None, - guard_cells=None, pml_cells=None, **kw): + @autoargs(exclude=['kw']) + def __init__(self, number_of_cells : picmi_types.VectorInt2 = None, + lower_bound : picmi_types.VectorFloat2 = None, + upper_bound : picmi_types.VectorFloat2 = None, + lower_boundary_conditions : picmi_types.VectorString2 = None, + upper_boundary_conditions : picmi_types.VectorString2 = None, + nr : int = None, + nz : int = None, + n_azimuthal_modes : int = None, + rmin : float = None, + rmax : float = None, + zmin : float = None, + zmax : float = None, + bc_rmin : str = None, + bc_rmax : str = None, + bc_zmin : str = None, + bc_zmax : str = None, + moving_window_velocity : picmi_types.VectorFloat2 = None, + refined_regions : Sequence[list] = None, + lower_bound_particles : picmi_types.VectorFloat2 = None, + upper_bound_particles : picmi_types.VectorFloat2 = None, + rmin_particles : float = None, + rmax_particles : float = None, + zmin_particles : float = None, + zmax_particles : float = None, + lower_boundary_conditions_particles : picmi_types.VectorString2 = None, + upper_boundary_conditions_particles : picmi_types.VectorString2 = None, + bc_rmin_particles : str = None, + bc_rmax_particles : str = None, + bc_zmin_particles : str = None, + bc_zmax_particles : str = None, + guard_cells : picmi_types.VectorInt2 = None, + pml_cells : picmi_types.VectorInt2 = None, + **kw): # Sanity check and init of input arguments related to grid parameters assert (number_of_cells is None) and (nr is not None and nz is not None) or \ @@ -357,84 +383,66 @@ def __init__(self, number_of_cells=None, lower_bound=None, upper_bound=None, Exception('Either upper_boundary_conditions or bc_rmax and bc_zmax must be specified') if number_of_cells is None: - number_of_cells = [nr, nz] + self.number_of_cells = [nr, nz] if lower_bound is None: - lower_bound = [rmin, zmin] + self.lower_bound = [rmin, zmin] if upper_bound is None: - upper_bound = [rmax, zmax] + self.upper_bound = [rmax, zmax] if lower_boundary_conditions is None: - lower_boundary_conditions = [bc_rmin, bc_zmin] + self.lower_boundary_conditions = [bc_rmin, bc_zmin] if upper_boundary_conditions is None: - upper_boundary_conditions = [bc_rmax, bc_zmax] + self.upper_boundary_conditions = [bc_rmax, bc_zmax] # Sanity check and init of input arguments related to particle boundary parameters # By default, if not specified, particle boundary values are the same as field boundary values # By default, if not specified, particle boundary conditions are the same as field boundary conditions if lower_bound_particles is None: if (rmin_particles is None) and (zmin_particles is None): - lower_bound_particles = lower_bound + self.lower_bound_particles = lower_bound else: - lower_bound_particles = [rmin_particles, zmin_particles] + self.lower_bound_particles = [rmin_particles, zmin_particles] if upper_bound_particles is None: if (rmax_particles is None) and (zmax_particles is None): - upper_bound_particles = upper_bound + self.upper_bound_particles = upper_bound else: - upper_bound_particles=[rmax_particles, zmax_particles] + self.upper_bound_particles = [rmax_particles, zmax_particles] if lower_boundary_conditions_particles is None: if (bc_rmin_particles is None) and (bc_zmin_particles is None): - lower_boundary_conditions_particles = lower_boundary_conditions + self.lower_boundary_conditions_particles = lower_boundary_conditions else: - lower_boundary_conditions_particles = [bc_rmin_particles, bc_zmin_particles] + self.lower_boundary_conditions_particles = [bc_rmin_particles, bc_zmin_particles] if upper_boundary_conditions_particles is None: if (bc_rmax_particles is None) and (bc_zmax_particles is None): - upper_boundary_conditions_particles = upper_boundary_conditions + self.upper_boundary_conditions_particles = upper_boundary_conditions else: - upper_boundary_conditions_particles = [bc_rmax_particles, bc_zmax_particles] - - # Sanity check on dimensionality of vector quantities - assert len(number_of_cells) == 2, Exception('Wrong number of cells specified') - assert len(lower_bound) == 2, Exception('Wrong number of lower bounds specified') - assert len(upper_bound) == 2, Exception('Wrong number of upper bounds specified') - assert len(lower_boundary_conditions) == 2, Exception('Wrong number of lower boundary conditions specified') - assert len(upper_boundary_conditions) == 2, Exception('Wrong number of upper boundary conditions specified') - - self.number_of_cells = number_of_cells - - self.lower_bound = lower_bound - self.upper_bound = upper_bound - self.lower_boundary_conditions = lower_boundary_conditions - self.upper_boundary_conditions = upper_boundary_conditions - self.lower_bound_particles = lower_bound_particles - self.upper_bound_particles = upper_bound_particles - self.lower_boundary_conditions_particles = lower_boundary_conditions_particles - self.upper_boundary_conditions_particles = upper_boundary_conditions_particles - self.guard_cells = guard_cells - self.pml_cells = pml_cells - - self.n_azimuthal_modes = n_azimuthal_modes - - self.moving_window_velocity = moving_window_velocity - - self.refined_regions = refined_regions - for region in self.refined_regions: - if len(region) == 3: - region.append([2,2]) - assert len(region[1]) == 2, Exception('The lo extent of the refined region must be a vector of length 2') - assert len(region[2]) == 2, Exception('The hi extent of the refined region must be a vector of length 2') - assert len(region[3]) == 2, Exception('The refinement factor of the refined region must be a vector of length 2') + self.upper_boundary_conditions_particles = [bc_rmax_particles, bc_zmax_particles] + + self.refined_regions = [] + if refined_regions is not None: + for region in refined_regions: + self.add_refined_region(*region) self.handle_init(kw) - def add_refined_region(self, level, lo, hi, refinement_factor=[2,2]): + def add_refined_region(self, level : int, + lo : picmi_types.VectorInt2, + hi : picmi_types.VectorInt2, + refinement_factor : picmi_types.VectorInt2 = None): """Add a refined region. - level: the refinement level, with 1 being the first level of refinement, 2 being the second etc. - lo, hi: vectors of length 2 specifying the extent of the region - refinement_factor: defaulting to [2,2] (relative to next lower level) + - level: the refinement level, with 1 being the first level of refinement, 2 being the second etc. + - lo, hi: vectors of length 2 specifying the extent of the region + - refinement_factor: defaulting to [2,2] (relative to next lower level) """ + if refinement_factor is None: + refinement_factor = [2,2] + assert len(lo) == 2, Exception('The lo extent of the refined region must be a vector of length 2') + assert len(hi) == 2, Exception('The hi extent of the refined region must be a vector of length 2') + assert len(refinement_factor) == 2, Exception('The refinement factor of the refined region must be a vector of length 2') self.refined_regions.append([level, lo, hi, refinement_factor]) +@typechecked class PICMI_Cartesian2DGrid(_ClassWithInit): """ Two dimensional Cartesian grid @@ -487,17 +495,39 @@ class PICMI_Cartesian2DGrid(_ClassWithInit): number_of_dimensions = 2 - def __init__(self, number_of_cells=None, lower_bound=None, upper_bound=None, - lower_boundary_conditions=None, upper_boundary_conditions=None, - nx=None, ny=None, - xmin=None, xmax=None, ymin=None, ymax=None, - bc_xmin=None, bc_xmax=None, bc_ymin=None, bc_ymax=None, - moving_window_velocity=None, refined_regions=[],lower_bound_particles=None, upper_bound_particles=None, - xmin_particles=None, xmax_particles=None, ymin_particles=None, ymax_particles=None, - lower_boundary_conditions_particles=None, upper_boundary_conditions_particles=None, - bc_xmin_particles=None, bc_xmax_particles=None, bc_ymin_particles=None, bc_ymax_particles=None, - guard_cells=None, pml_cells=None, - **kw): + @autoargs(exclude=['kw']) + def __init__(self, number_of_cells : picmi_types.VectorInt2 = None, + lower_bound : picmi_types.VectorFloat2 = None, + upper_bound : picmi_types.VectorFloat2 = None, + lower_boundary_conditions : picmi_types.VectorString2 = None, + upper_boundary_conditions : picmi_types.VectorString2 = None, + nx : int = None, + ny : int = None, + xmin : float = None, + xmax : float = None, + ymin : float = None, + ymax : float = None, + bc_xmin : str = None, + bc_xmax : str = None, + bc_ymin : str = None, + bc_ymax : str = None, + moving_window_velocity : picmi_types.VectorFloat2 = None, + refined_regions : Sequence[list] = None, + lower_bound_particles : picmi_types.VectorFloat2 = None, + upper_bound_particles : picmi_types.VectorFloat2 = None, + xmin_particles : float = None, + xmax_particles : float = None, + ymin_particles : float = None, + ymax_particles : float = None, + lower_boundary_conditions_particles : picmi_types.VectorString2 = None, + upper_boundary_conditions_particles : picmi_types.VectorString2 = None, + bc_xmin_particles : str = None, + bc_xmax_particles : str = None, + bc_ymin_particles : str = None, + bc_ymax_particles : str = None, + guard_cells : picmi_types.VectorInt2 = None, + pml_cells : picmi_types.VectorInt2 = None, + **kw): # Sanity check and init of input arguments related to grid parameters assert (number_of_cells is None) and (nx is not None and ny is not None) or \ @@ -517,86 +547,66 @@ def __init__(self, number_of_cells=None, lower_bound=None, upper_bound=None, Exception('Either upper_boundary_conditions or bc_xmax and bc_ymax must be specified') if number_of_cells is None: - number_of_cells = [nx, ny] + self.number_of_cells = [nx, ny] if lower_bound is None: - lower_bound = [xmin, ymin] + self.lower_bound = [xmin, ymin] if upper_bound is None: - upper_bound = [xmax, ymax] + self.upper_bound = [xmax, ymax] if lower_boundary_conditions is None: - lower_boundary_conditions = [bc_xmin, bc_ymin] + self.lower_boundary_conditions = [bc_xmin, bc_ymin] if upper_boundary_conditions is None: - upper_boundary_conditions = [bc_xmax, bc_ymax] + self.upper_boundary_conditions = [bc_xmax, bc_ymax] # Sanity check and init of input arguments related to particle boundary parameters # By default, if not specified, particle boundary values are the same as field boundary values # By default, if not specified, particle boundary conditions are the same as field boundary conditions if lower_bound_particles is None: if (xmin_particles is None) and (ymin_particles is None): - lower_bound_particles = lower_bound + self.lower_bound_particles = lower_bound else: - lower_bound_particles = [xmin_particles, ymin_particles] + self.lower_bound_particles = [xmin_particles, ymin_particles] if upper_bound_particles is None: if (xmax_particles is None) and (ymax_particles is None): - upper_bound_particles = upper_bound + self.upper_bound_particles = upper_bound else: - upper_bound_particles=[xmax_particles, ymax_particles] + self.upper_bound_particles = [xmax_particles, ymax_particles] if lower_boundary_conditions_particles is None: if (bc_xmin_particles is None) and (bc_ymin_particles is None): - lower_boundary_conditions_particles = lower_boundary_conditions + self.lower_boundary_conditions_particles = lower_boundary_conditions else: - lower_boundary_conditions_particles = [bc_xmin_particles, bc_ymin_particles] + self.lower_boundary_conditions_particles = [bc_xmin_particles, bc_ymin_particles] if upper_boundary_conditions_particles is None: if (bc_xmax_particles is None) and (bc_ymax_particles is None): - upper_boundary_conditions_particles = upper_boundary_conditions + self.upper_boundary_conditions_particles = upper_boundary_conditions else: - upper_boundary_conditions_particles = [bc_xmax_particles, bc_ymax_particles] - - # Sanity check on dimensionality of vector quantities - assert len(number_of_cells) == 2, Exception('Wrong number of cells specified') - assert len(lower_bound) == 2, Exception('Wrong number of lower bounds specified') - assert len(upper_bound) == 2, Exception('Wrong number of upper bounds specified') - assert len(lower_boundary_conditions) == 2, Exception('Wrong number of lower boundary conditions specified') - assert len(upper_boundary_conditions) == 2, Exception('Wrong number of upper boundary conditions specified') - assert len(lower_bound_particles) == 2, Exception('Wrong number of particle lower bounds specified') - assert len(upper_bound_particles) == 2, Exception('Wrong number of particle upper bounds specified') - assert len(lower_boundary_conditions_particles) == 2, Exception('Wrong number of lower particle boundary conditions specified') - assert len(upper_boundary_conditions_particles) == 2, Exception('Wrong number of upper particle boundary conditions specified') - - self.number_of_cells = number_of_cells - self.lower_bound = lower_bound - self.upper_bound = upper_bound - self.lower_boundary_conditions = lower_boundary_conditions - self.upper_boundary_conditions = upper_boundary_conditions - self.lower_bound_particles = lower_bound_particles - self.upper_bound_particles = upper_bound_particles - self.lower_boundary_conditions_particles = lower_boundary_conditions_particles - self.upper_boundary_conditions_particles = upper_boundary_conditions_particles - self.guard_cells = guard_cells - self.pml_cells = pml_cells - - self.moving_window_velocity = moving_window_velocity - - self.refined_regions = refined_regions - - for region in self.refined_regions: - if len(region) == 3: - region.append([2,2]) - assert len(region[1]) == 2, Exception('The lo extent of the refined region must be a vector of length 2') - assert len(region[2]) == 2, Exception('The hi extent of the refined region must be a vector of length 2') - assert len(region[3]) == 2, Exception('The refinement factor of the refined region must be a vector of length 2') + self.upper_boundary_conditions_particles = [bc_xmax_particles, bc_ymax_particles] + + self.refined_regions = [] + if refined_regions is not None: + for region in refined_regions: + self.add_refined_region(*region) self.handle_init(kw) - def add_refined_region(self, level, lo, hi, refinement_factor=[2,2]): + def add_refined_region(self, level : int, + lo : picmi_types.VectorInt2, + hi : picmi_types.VectorInt2, + refinement_factor : picmi_types.VectorInt2 = None): """Add a refined region. - level: the refinement level, with 1 being the first level of refinement, 2 being the second etc. - lo, hi: vectors of length 2 specifying the extent of the region - refinement_factor: defaulting to [2,2] (relative to next lower level) + - level: the refinement level, with 1 being the first level of refinement, 2 being the second etc. + - lo, hi: vectors of length 2 specifying the extent of the region + - refinement_factor: defaulting to [2,2] (relative to next lower level) """ + if refinement_factor is None: + refinement_factor = [2,2] + assert len(lo) == 2, Exception('The lo extent of the refined region must be a vector of length 2') + assert len(hi) == 2, Exception('The hi extent of the refined region must be a vector of length 2') + assert len(refinement_factor) == 2, Exception('The refinement factor of the refined region must be a vector of length 2') self.refined_regions.append([level, lo, hi, refinement_factor]) +@typechecked class PICMI_Cartesian3DGrid(_ClassWithInit): """ Three dimensional Cartesian grid @@ -657,17 +667,48 @@ class PICMI_Cartesian3DGrid(_ClassWithInit): number_of_dimensions = 3 - def __init__(self, number_of_cells=None, lower_bound=None, upper_bound=None, - lower_boundary_conditions=None, upper_boundary_conditions=None, - nx=None, ny=None, nz=None, - xmin=None, xmax=None, ymin=None, ymax=None, zmin=None, zmax=None, - bc_xmin=None, bc_xmax=None, bc_ymin=None, bc_ymax=None, bc_zmin=None, bc_zmax=None, - moving_window_velocity=None, refined_regions=[], lower_bound_particles=None, upper_bound_particles=None, - xmin_particles=None, xmax_particles=None, ymin_particles=None, ymax_particles=None, zmin_particles=None, zmax_particles=None, - lower_boundary_conditions_particles=None, upper_boundary_conditions_particles=None, - bc_xmin_particles=None, bc_xmax_particles=None, bc_ymin_particles=None, bc_ymax_particles=None, - bc_zmin_particles=None, bc_zmax_particles=None, guard_cells=None, pml_cells=None, - **kw): + @autoargs(exclude=['kw']) + def __init__(self, number_of_cells : picmi_types.VectorInt3 = None, + lower_bound : picmi_types.VectorFloat3 = None, + upper_bound : picmi_types.VectorFloat3 = None, + lower_boundary_conditions : picmi_types.VectorString3 = None, + upper_boundary_conditions : picmi_types.VectorString3 = None, + nx : int = None, + ny : int = None, + nz : int = None, + xmin : float = None, + xmax : float = None, + ymin : float = None, + ymax : float = None, + zmin : float = None, + zmax : float = None, + bc_xmin : str = None, + bc_xmax : str = None, + bc_ymin : str = None, + bc_ymax : str = None, + bc_zmin : str = None, + bc_zmax : str = None, + moving_window_velocity : picmi_types.VectorFloat3 = None, + refined_regions : Sequence[list] = None, + lower_bound_particles : picmi_types.VectorFloat3 = None, + upper_bound_particles : picmi_types.VectorFloat3 = None, + xmin_particles : float = None, + xmax_particles : float = None, + ymin_particles : float = None, + ymax_particles : float = None, + zmin_particles : float = None, + zmax_particles : float = None, + lower_boundary_conditions_particles : picmi_types.VectorString3 = None, + upper_boundary_conditions_particles : picmi_types.VectorString3 = None, + bc_xmin_particles : str = None, + bc_xmax_particles : str = None, + bc_ymin_particles : str = None, + bc_ymax_particles : str = None, + bc_zmin_particles : str = None, + bc_zmax_particles : str = None, + guard_cells : picmi_types.VectorInt3 = None, + pml_cells : picmi_types.VectorInt3 = None, + **kw): # Sanity check and init of input arguments related to grid parameters assert (number_of_cells is None) and (nx is not None and ny is not None and nz is not None) or \ @@ -687,80 +728,60 @@ def __init__(self, number_of_cells=None, lower_bound=None, upper_bound=None, Exception('Either upper_boundary_conditions or bc_xmax, bc_ymax, and bc_zmax must be specified') if number_of_cells is None: - number_of_cells = [nx, ny, nz] + self.number_of_cells = [nx, ny, nz] if lower_bound is None: - lower_bound = [xmin, ymin, zmin] + self.lower_bound = [xmin, ymin, zmin] if upper_bound is None: - upper_bound = [xmax, ymax, zmax] + self.upper_bound = [xmax, ymax, zmax] if lower_boundary_conditions is None: - lower_boundary_conditions = [bc_xmin, bc_ymin, bc_zmin] + self.lower_boundary_conditions = [bc_xmin, bc_ymin, bc_zmin] if upper_boundary_conditions is None: - upper_boundary_conditions = [bc_xmax, bc_ymax, bc_zmax] + self.upper_boundary_conditions = [bc_xmax, bc_ymax, bc_zmax] # Sanity check and init of input arguments related to particle boundary parameters # By default, if not specified, particle boundary values are the same as field boundary values # By default, if not specified, particle boundary conditions are the same as field boundary conditions if lower_bound_particles is None: if (xmin_particles is None) and (ymin_particles is None) and (zmin_particles is None): - lower_bound_particles = lower_bound + self.lower_bound_particles = lower_bound else: - lower_bound_particles = [xmin_particles, ymin_particles, zmin_particles] + self.lower_bound_particles = [xmin_particles, ymin_particles, zmin_particles] if upper_bound_particles is None: if (xmax_particles is None) and (ymax_particles is None) and (zmax_particles is None): - upper_bound_particles = upper_bound + self.upper_bound_particles = upper_bound else: - upper_bound_particles = [xmax_particles, ymax_particles, zmax_particles] + self.upper_bound_particles = [xmax_particles, ymax_particles, zmax_particles] if lower_boundary_conditions_particles is None: if (bc_xmin_particles is None) and (bc_ymin_particles is None) and (bc_zmin_particles is None): - lower_boundary_conditions_particles = lower_boundary_conditions + self.lower_boundary_conditions_particles = lower_boundary_conditions else: - lower_boundary_conditions_particles = [bc_xmin_particles, bc_ymin_particles, bc_zmin_particles] + self.lower_boundary_conditions_particles = [bc_xmin_particles, bc_ymin_particles, bc_zmin_particles] if upper_boundary_conditions_particles is None: if (bc_xmax_particles is None) and (bc_ymax_particles is None) and (bc_zmax_particles is None): - upper_boundary_conditions_particles = upper_boundary_conditions + self.upper_boundary_conditions_particles = upper_boundary_conditions else: - upper_boundary_conditions_particles = [bc_xmax_particles, bc_ymax_particles, bc_zmax_particles] - - # Sanity check on number of arguments of vector quantities - assert len(number_of_cells) == 3, Exception('Wrong number of cells specified') - assert len(lower_bound) == 3, Exception('Wrong number of lower bounds specified') - assert len(upper_bound) == 3, Exception('Wrong number of upper bounds specified') - assert len(lower_boundary_conditions) == 3, Exception('Wrong number of lower boundary conditions specified') - assert len(upper_boundary_conditions) == 3, Exception('Wrong number of upper boundary conditions specified') - assert len(lower_bound_particles) == 3, Exception('Wrong number of particle lower bounds specified') - assert len(upper_bound_particles) == 3, Exception('Wrong number of particle upper bounds specified') - assert len(lower_boundary_conditions_particles) == 3, Exception('Wrong number of particle lower boundary conditions specified') - assert len(upper_boundary_conditions_particles) == 3, Exception('Wrong number of particle upper boundary conditions specified') - - self.number_of_cells = number_of_cells - self.lower_bound = lower_bound - self.upper_bound = upper_bound - self.lower_boundary_conditions = lower_boundary_conditions - self.upper_boundary_conditions = upper_boundary_conditions - self.lower_bound_particles = lower_bound_particles - self.upper_bound_particles = upper_bound_particles - self.lower_boundary_conditions_particles = lower_boundary_conditions_particles - self.upper_boundary_conditions_particles = upper_boundary_conditions_particles - self.guard_cells = guard_cells - self.pml_cells = pml_cells - - self.moving_window_velocity = moving_window_velocity - - self.refined_regions = refined_regions - for region in self.refined_regions: - if len(region) == 3: - region.append([2,2,2]) - assert len(region[1]) == 3, Exception('The lo extent of the refined region must be a vector of length 3') - assert len(region[2]) == 3, Exception('The hi extent of the refined region must be a vector of length 3') - assert len(region[3]) == 3, Exception('The refinement factor of the refined region must be a vector of length 3') + self.upper_boundary_conditions_particles = [bc_xmax_particles, bc_ymax_particles, bc_zmax_particles] + + self.refined_regions = [] + if refined_regions is not None: + for region in refined_regions: + self.add_refined_region(*region) self.handle_init(kw) - def add_refined_region(self, level, lo, hi, refinement_factor=[2,2,2]): + def add_refined_region(self, level : int, + lo : picmi_types.VectorInt3, + hi : picmi_types.VectorInt3, + refinement_factor : picmi_types.VectorInt3 = None): """Add a refined region. - level: the refinement level, with 1 being the first level of refinement, 2 being the second etc. - lo, hi: vectors of length 3 specifying the extent of the region - refinement_factor: defaulting to [2,2,2] (relative to next lower level) + - level: the refinement level, with 1 being the first level of refinement, 2 being the second etc. + - lo, hi: vectors of length 3 specifying the extent of the region + - refinement_factor: defaulting to [2,2,2] (relative to next lower level) """ + if refinement_factor is None: + refinement_factor = [2,2,2] + assert len(lo) == 3, Exception('The lo extent of the refined region must be a vector of length 3') + assert len(hi) == 3, Exception('The hi extent of the refined region must be a vector of length 3') + assert len(refinement_factor) == 3, Exception('The refinement factor of the refined region must be a vector of length 3') self.refined_regions.append([level, lo, hi, refinement_factor]) diff --git a/PICMI_Python/picmi_types.py b/PICMI_Python/picmi_types.py index 2e56d3d..54166d7 100644 --- a/PICMI_Python/picmi_types.py +++ b/PICMI_Python/picmi_types.py @@ -5,9 +5,19 @@ import numpy as np +VectorFloat1 = typing.NewType('VectorFloat1', typing.Union[Sequence[float], np.ndarray[1, np.float64]]) +VectorInt1 = typing.NewType('VectorInt1', typing.Union[Sequence[int], np.ndarray[1, np.int64]]) +VectorString1 = typing.NewType('VectorString1', Sequence[str]) + +VectorFloat2 = typing.NewType('VectorFloat2', typing.Union[Sequence[float], np.ndarray[2, np.float64]]) +VectorInt2 = typing.NewType('VectorInt2', typing.Union[Sequence[int], np.ndarray[2, np.int64]]) +VectorString2 = typing.NewType('VectorString2', Sequence[str]) + VectorFloat3 = typing.NewType('VectorFloat3', typing.Union[Sequence[float], np.ndarray[3, np.float64]]) -VectorFloatNone3 = typing.NewType('VectorFloat3', typing.Union[Sequence[typing.Union[float, None]], np.ndarray[3, np.float64]]) +VectorFloatNone3 = typing.NewType('VectorFloatNone3', typing.Union[Sequence[typing.Union[float, None]], np.ndarray[3, np.float64]]) VectorInt3 = typing.NewType('VectorInt3', typing.Union[Sequence[int], np.ndarray[3, np.int64]]) +VectorBool3 = typing.NewType('VectorBool3', Sequence[bool]) +VectorString3 = typing.NewType('VectorString3', Sequence[str]) Expression = typing.NewType('Expression', str) VectorExpression3 = typing.NewType('VectorExpression3', Sequence[str]) @@ -17,6 +27,13 @@ fields.PICMI_CylindricalGrid, fields.PICMI_Cartesian2DGrid, fields.PICMI_Cartesian3DGrid])""" +# This is needed in the fields.py file since it cannot import itself +FieldsGridType = """typing.NewType('GridType', typing.Union[PICMI_Cartesian1DGrid, + PICMI_CylindricalGrid, + PICMI_Cartesian2DGrid, + PICMI_Cartesian3DGrid])""" + +SmootherType = """typing.NewType('SmootherType', PICMI_BinomialSmoother)""" _species_union = "typing.Union[particles.PICMI_Species, particles.PICMI_MultiSpecies]" From 28795a474d1c7dfd7b210cea8a6351235b41d793 Mon Sep 17 00:00:00 2001 From: Dave Grote Date: Mon, 11 Apr 2022 14:30:20 -0700 Subject: [PATCH 12/13] Small fix of n_physical_particles type --- PICMI_Python/particles.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/PICMI_Python/particles.py b/PICMI_Python/particles.py index ce3f1c7..333ad1c 100644 --- a/PICMI_Python/particles.py +++ b/PICMI_Python/particles.py @@ -170,7 +170,7 @@ class PICMI_GaussianBunchDistribution(_ClassWithInit): - velocity_divergence=[0,0,0]: Expansion rate of the bunch at t=0 (vector) [m/s/m] """ @autoargs(exclude=['kw']) - def __init__(self,n_physical_particles : int, + def __init__(self,n_physical_particles : float, rms_bunch_size : picmi_types.VectorFloat3, rms_velocity : picmi_types.VectorFloat3 = [0.,0.,0.], centroid_position : picmi_types.VectorFloat3 = [0.,0.,0.], From da7a09bb5790e6969b00b41a71928778853c4992 Mon Sep 17 00:00:00 2001 From: Dave Grote Date: Mon, 11 Apr 2022 14:47:57 -0700 Subject: [PATCH 13/13] Small fixes in simulation.py --- PICMI_Python/simulation.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/PICMI_Python/simulation.py b/PICMI_Python/simulation.py index 23a2e1b..13901b3 100644 --- a/PICMI_Python/simulation.py +++ b/PICMI_Python/simulation.py @@ -15,6 +15,7 @@ from . import lasers from . import diagnostics from . import picmi_types +from . import applied_fields # --------------------- # Main simulation object @@ -86,7 +87,7 @@ def __init__(self, solver : picmi_types.SolverType = None, self.handle_init(kw) def add_species(self, species : picmi_types.SpeciesType, - layout : picmi_types.LayoutType, + layout : picmi_types.LayoutType = None, initialize_self_field : bool = None): """ Add species to be used in the simulation