work on poincare seeds

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")

w7x/core.py

-"""
-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"