|
1 | 1 | """ |
2 | 2 | This file is part of PIConGPU. |
3 | | -Copyright 2021-2024 PIConGPU contributors |
4 | | -Authors: Hannes Troepgen, Brian Edward Marre |
| 3 | +Copyright 2021-2025 PIConGPU contributors |
| 4 | +Authors: Hannes Troepgen, Brian Edward Marre, Julian Lenz |
5 | 5 | License: GPLv3+ |
6 | 6 | """ |
7 | 7 |
|
| 8 | +from numpy import vectorize |
8 | 9 | from ...pypicongpu import species |
9 | | -from ...pypicongpu import util |
10 | 10 |
|
11 | | -import picmistandard |
12 | 11 |
|
| 12 | +import logging |
13 | 13 | import typeguard |
14 | 14 | import typing |
15 | 15 | import sympy |
16 | | -import math |
| 16 | +import traceback |
17 | 17 |
|
18 | 18 | """ |
19 | 19 | note on rms_velocity: |
|
38 | 38 |
|
39 | 39 |
|
40 | 40 | @typeguard.typechecked |
41 | | -class AnalyticDistribution(picmistandard.PICMI_AnalyticDistribution): |
42 | | - """Analytic Particle Distribution as defined by PICMI @todo""" |
| 41 | +class AnalyticDistribution: |
| 42 | + """ |
| 43 | + This class represents a plasma with a density defined by an analytic expression. |
| 44 | +
|
| 45 | + The function must be constructed using sympy functions |
| 46 | + to enable code generation and manipulation. |
| 47 | + This is a slight deviation from the PICMI standard. |
| 48 | + Furthermore, we don't implement substitution of variables |
| 49 | + as suggested in the PICMI standard. |
| 50 | + Instead we propose that you write your function |
| 51 | + with further variables as keyword arguments |
| 52 | + and substitute them yourself before handing it over to AnalyticDistribution. |
| 53 | + See the end-to-end tests for examples of this. |
| 54 | +
|
| 55 | + Writing such functions (or rather writing sympy in general) |
| 56 | + comes with a few pitfalls but also advantages as listed below. |
| 57 | + Make sure that you familiarise yourself with writing sympy. |
| 58 | +
|
| 59 | + Advantages: |
| 60 | + - The sympy language is closer to mathematical language than to coding |
| 61 | + which might make it more natural to use for some physicists. |
| 62 | + - You can extract the exact distribution |
| 63 | + that was used from the member `density_expression` |
| 64 | + and use it any way you'd use any sympy expression. |
| 65 | + In particular, you can print it to various formats, |
| 66 | + say, LaTeX for automated inclusion in papers. |
| 67 | + - We can easily evaluate it from within python. |
| 68 | + This is what the __call__ operator does. |
| 69 | + However, code generation here can have some difficulties |
| 70 | + with advanced broadcasting for numpy variables. |
| 71 | + The operator implements a fallback in such cases. |
| 72 | + This fallback might be slightly slower on large inputs. |
| 73 | + Try rewriting your function, in case you experience performance problems. |
| 74 | +
|
| 75 | + Pitfalls: |
| 76 | + - We don't handle vectors yet. |
| 77 | + But that's probably not too important for density expressions. |
| 78 | + Just be explicit handling multiple vector components for now. |
| 79 | + - Some operations might compile to suboptimal code |
| 80 | + concerning numerical performance and stability. |
| 81 | + Experts might want to inspect the generated C++ code |
| 82 | + and/or check the unit tests for the PMAccPrinter |
| 83 | + to find the precise mapping of sympy expressions |
| 84 | + to PMAcc code. |
| 85 | + Please approach us if you should stumble across this. |
| 86 | + - Control flow is an interesting topic in this regard. |
| 87 | + With respect to your three position coordinates |
| 88 | + (which will be sympy symbols internally), |
| 89 | + you must use pure sympy, e.g., |
| 90 | + replacing if-conditions with sympy.Piecewise and so on. |
| 91 | + With respect to further parameters, |
| 92 | + you're free to use any python construct you want |
| 93 | + (if-conditions, loops, etc.). |
| 94 | + - sympy.Piecewise has the potentially surprising property |
| 95 | + that any pieces that you leave undefined are interpreted as nan. |
| 96 | + This implies that adding two complementary sympy.Piecewise |
| 97 | + renders the whole expression nan and not -- as you might expect -- |
| 98 | + defined on the union of the defined regions. |
| 99 | + There are two options to circumvent this: |
| 100 | + Either you can define multiple sympy.Piecewise with |
| 101 | + (0.0, True) as the last condition which means 0 everywhere else. |
| 102 | + (Make sure it's the last!) |
| 103 | + Summing those up, works just as you'd expect. |
| 104 | + Alternatively, you can define only the (expression, condition) tuples |
| 105 | + and assemble them in a sympy.Piecewise in one go. |
| 106 | + That's the way chosen in the end-to-end tests. |
| 107 | +
|
| 108 | + Parameters: |
| 109 | + density_expression (Callable): |
| 110 | + A Python function that takes x, y, z coordinates (in SI units) |
| 111 | + and returns the density (in SI units) at that point. |
| 112 | + It should use sympy functionality. |
| 113 | + directed_velocity (3-tuple of float): |
| 114 | + A collective velocity for the particle distribution. |
| 115 | + (currently untested) |
| 116 | + """ |
| 117 | + |
| 118 | + def __init__(self, density_expression, directed_velocity=(0.0, 0.0, 0.0)): |
| 119 | + self.density_function = density_expression |
| 120 | + self.rms_velocity = (0.0, 0.0, 0.0) |
| 121 | + self.directed_velocity = tuple(float(v) for v in directed_velocity) |
| 122 | + x, y, z = sympy.symbols("x,y,z") |
| 123 | + self.density_expression = ( |
| 124 | + # We add the simplify because of the following: |
| 125 | + # Translating to C++ requires ALL cases of a Piecewise to be defined |
| 126 | + # such that there's a fallback for if-conditions. |
| 127 | + # The user might have written their formula piecing together |
| 128 | + # multiple partial Piecewise instances. |
| 129 | + # Without simplification, sympy tries to translate them individually. |
| 130 | + # and fails to do so even if they supplement each other |
| 131 | + # into a function defined everywhere. |
| 132 | + sympy.simplify(self.density_function(x, y, z)) |
| 133 | + # density_expression might be independent of any or all of the three variables. |
| 134 | + # (This might even happen due to the simplification.) |
| 135 | + # In order to be sure to arrive at a function of these three variables, |
| 136 | + # we add this trivial additional term. |
| 137 | + + (0 * x * y * z) |
| 138 | + ) |
| 139 | + self.warned_about_lambdify_failure = False |
| 140 | + |
| 141 | + def get_as_pypicongpu(self, grid) -> species.operation.densityprofile.DensityProfile: |
| 142 | + x, y, z = sympy.symbols("x,y,z") |
| 143 | + return species.operation.densityprofile.FreeFormula(density_expression=self.density_expression) |
43 | 144 |
|
44 | 145 | def picongpu_get_rms_velocity_si(self) -> typing.Tuple[float, float, float]: |
45 | | - return tuple(self.rms_velocity) |
46 | | - |
47 | | - def get_as_pypicongpu(self) -> species.operation.densityprofile.DensityProfile: |
48 | | - util.unsupported("momentum expressions", self.momentum_expressions) |
49 | | - util.unsupported("fill in", self.fill_in) |
50 | | - |
51 | | - # TODO |
52 | | - profile = object() |
53 | | - profile.lower_bound = tuple(map(lambda x: -math.inf if x is None else x, self.lower_bound)) |
54 | | - profile.upper_bound = tuple(map(lambda x: math.inf if x is None else x, self.upper_bound)) |
55 | | - |
56 | | - # final (more thorough) formula checking will be invoked inside |
57 | | - # pypicongpu on translation to CPP |
58 | | - sympy_density_expression = sympy.sympify(self.density_expression).subs(self.user_defined_kw) |
59 | | - profile.expression = sympy_density_expression |
60 | | - |
61 | | - return profile |
| 146 | + return self.rms_velocity |
62 | 147 |
|
63 | 148 | def get_picongpu_drift(self) -> typing.Optional[species.operation.momentum.Drift]: |
64 | 149 | """ |
65 | 150 | Get drift for pypicongpu |
66 | 151 | :return: pypicongpu drift object or None |
67 | 152 | """ |
68 | | - if [0, 0, 0] == self.directed_velocity: |
| 153 | + if all(v == 0 for v in self.directed_velocity): |
69 | 154 | return None |
70 | 155 |
|
71 | 156 | drift = species.operation.momentum.Drift() |
72 | | - drift.fill_from_velocity(tuple(self.directed_velocity)) |
| 157 | + drift.fill_from_velocity(self.directed_velocity) |
73 | 158 | return drift |
| 159 | + |
| 160 | + def __call__(self, *args, **kwargs): |
| 161 | + try: |
| 162 | + # This produces faster code but the code generation is not perfect. |
| 163 | + # There are cases where the generated code can't handle broadcasting properly. |
| 164 | + return sympy.lambdify(sympy.symbols("x,y,z"), self.density_expression, "numpy")(*args, **kwargs) |
| 165 | + except ValueError: |
| 166 | + if not self.warned_about_lambdify_failure: |
| 167 | + message = ( |
| 168 | + "Sympy did not manage to produce proper numpy code for your AnalyticDistribution. " |
| 169 | + "If you run into performance problems, try to rewrite your function. " |
| 170 | + "Here's the original error message:" |
| 171 | + ) |
| 172 | + logging.warning(message) |
| 173 | + logging.warning(traceback.format_exc()) |
| 174 | + logging.warning("Continuing operation using a slower serialised version now.") |
| 175 | + self.warned_about_lambdify_failure = True |
| 176 | + # This basically calls the original function in a big loop. |
| 177 | + # Slower but more reliable in some cases of difficult broadcasting. |
| 178 | + return vectorize(self.density_function)(*args, **kwargs) |
0 commit comments