diff --git a/examples/poincare.py b/examples/poincare.py index 03d84ffdf8dabb5b47df34d68edaa68666c6cc33..573b6f2c182978e08eaaa7d75d03d7ce833b1bdc 100644 --- a/examples/poincare.py +++ b/examples/poincare.py @@ -7,8 +7,12 @@ Author: daniel.boeckenhoff@ipp.mpg.de ################################################## import logging +import argparse import numpy as np +import ast +from numpy.lib.arraysetops import isin import rna +import rna.parsing import tfields import w7x import w7x.simulation.flt @@ -17,26 +21,90 @@ import w7x.simulation.components logging.basicConfig(level=logging.INFO) +# pylint:disable=protected-access +class PhiAction(argparse._StoreAction): + def __call__(self, parser, namespace, values, option_string=None): + super().__call__(parser, namespace, values, option_string=option_string) + for i, phi in enumerate(namespace.phi): + if "*pi" in phi: + phi = phi.replace("*pi", "") + phi_mul = np.pi + else: + phi_mul = 1 + namespace.phi[i] = phi_mul * ast.literal_eval(phi) + import pdb + + pdb.set_trace() + # namespace phi = np.array(args.phi) / 180 * np.pi + + +# pylint:disable=protected-access +class SeedAction(argparse._StoreAction): + def __call__(self, parser, namespace, values, option_string=None): + super().__call__(parser, namespace, values, option_string=option_string) + for i, seed in enumerate(namespace.seeds): + import pdb + + pdb.set_trace() + + +class PoincareParser(argparse.ArgumentParser): + """ + Parent parser with the "level" argument, automatically setting the logging level + """ + + def __init__(self, add_help=False, **kwargs): + super().__init__(add_help=add_help, **kwargs) + self.add_argument( + "--phi", + nargs="*", + default=[0.0], + action=PhiAction, + help="Torodial angle(s) [rad]. Multiple angles possible. You can use '*pi' to convert from deg.", + ) + self.add_argument( + "--seeds", + nargs="*", + action=SeedAction, + default=tfields.Points3D([[5.8, 0.0, 0.0], [6.0, 0.0, 0.0]]), + help="Poincare seeds. Multiple ways of passing seeds. str, comma separated list,...", + ) + self.add_argument( + "--n_points", + type=int, + default=300, + help="Number of points to trace one seed.", + ) + + def main(): """poincare plot creation""" - phi = 171.2 / 180 * np.pi + parser = argparse.ArgumentParser( + parents=[rna.parsing.LogParser(), PoincareParser()] + ) + parser.add_argument( + "--distribute", "-d", action="store_true", help="distribute the work load" + ) + args = parser.parse_args() + assembly = w7x.model.Assembly( components=[ w7x.config.AssemblyGroups.Vessel(), w7x.config.AssemblyGroups.Divertor(), + w7x.config.AssemblyGroups.Baffle(), ] ) flt = w7x.simulation.flt.FieldLineTracer() components = w7x.simulation.components.Components() - with w7x.distribute(True): + with w7x.distribute(args.distribute): graph = w7x.State.merged( components.mesh_slice(assembly, phi=phi), # pylint:disable=no-value-for-parameter flt.trace_poincare( phi=phi, - seeds=tfields.Points3D([[5.8, 0.0, 0.0], [6.0, 0.0, 0.0]]), - n_points=100, + seeds=args.seeds, + n_points=args.n_points, ), ) graph.visualize(filename="poincare_graph.pdf") diff --git a/w7x/core.py b/w7x/core.py index 571429a05b8020366ca6deae61d15d98739dabcd..364383a8b48c3cbfca866be1b586e0a67c66f806 100644 --- a/w7x/core.py +++ b/w7x/core.py @@ -1,13 +1,3 @@ -""" -Zen: - * Backend methods return primitives - * Backend methods do not mutate -> dask.delayed - -TODO-2: Question: how to support optional dependencies (e.g. FLT << Equilibrium | Biot-savart) - -> maybe not necessary if Equilibrium returns None for 0 pressure - -> BranchPythonOperator? airflow/example_dags/example_branch_python_dop_operator_3.py -""" - import abc import typing import logging @@ -34,6 +24,9 @@ class Backend(rna.pattern.backend.Backend): Adaptaion of the 'Strategy' class from the strategy design pattern. Backends are the meat of the implementation for a certain variant of a code. + Zen: + * Backend methods return primitives + * Backend methods do not mutate -> dask.delayed """ STRATEGY_MODULE_BASE = "w7x.simulation.backends"