From 2b96707f3147964390ef3339626d5cb3d1aa27fe Mon Sep 17 00:00:00 2001 From: emile grivet Date: Thu, 7 May 2026 14:06:53 +0200 Subject: [PATCH 1/5] preparation of branch --- src/struphy/kinetic_background/base.py | 225 ++++++++++++++++++ src/struphy/kinetic_background/maxwellians.py | 35 +++ .../kinetic_background/tests/test_base.py | 55 ++++- 3 files changed, 314 insertions(+), 1 deletion(-) diff --git a/src/struphy/kinetic_background/base.py b/src/struphy/kinetic_background/base.py index a223ab7d6..940163eb2 100644 --- a/src/struphy/kinetic_background/base.py +++ b/src/struphy/kinetic_background/base.py @@ -4,9 +4,14 @@ from typing import Callable import cunumpy as xp +import matplotlib.pyplot as plt +from matplotlib import cm +from matplotlib.colors import Normalize from struphy.fields_background.base import FluidEquilibriumWithB +from struphy.geometry.base import Domain from struphy.initial.base import Perturbation +from struphy.io.options import LiteralOptions class KineticBackground(metaclass=ABCMeta): @@ -109,6 +114,226 @@ def __call__(self, *args): """ pass + def plot_density_profile( + self, + dim_1: LiteralOptions.DimensionToPlot = "e1", + dim_2: LiteralOptions.DimensionToPlot | None = None, + v_lim: float = 5.0, + resol: int = 100, + integrate_resol: int = 50, + logical_coord: tuple[float] = (0.5, 0.5, 0.5), + in_physical: bool = False, + domain: Domain | None = None, + proj_axis: tuple[float,] = (0, 1), + plot_3D: bool = False, + title: str | None = None, + **kwargs, + ): + """ + Plots the density profile of an initial condition (or a background) from the phase space distribution. The projection can either be 1D or 2D, in the logical space or in cartesian. + + Parameters + ---------- + dim_1, dim_2 : LiteralOptions = ["e1","e2","e3","v1","v2","v3"] + The axes used in the projection, they refere to logical space axes. If dim_2 is not defined the projection is 1D, it is 2D if dim_2 is attributed. + + v_lim : float = 5.0 + Limit value of the velocity axes. + + resol : int = 100 + Resolution along each axes + + integrate_resol : int = 10 + Resolution along not used velocity axes. The density is reduced (with a maximum function) over these axes before being plotted. High values (>50) may require much memory. + + logical_coord : tuple[float] = (0.5, 0.5, 0.5) + Refere to the default coordinate (in logical space) attributed to each axe which is not used in the projection. + + in_physical : bool = False + Specify if the result is plotted in logical coordinates or in cartesian coordinates, has a effect in 2D plotting. If True, you must specify a domain. + + domain : Domain | None = None + Domain used to plot the density if in_physical=True. + + proj_axis : tuple[float] = (0,1) + Axes of the cartesian coordinates used to plot the density: 0->x, 1->y, 2->z. + I you do not see the density profile in 2D, you may change these axes. + + plot_3D : bool = False + Plots a surface in a 3D environment. Only for physical projection. + """ + if "use_mu" in kwargs: + use_mu = kwargs["use_mu"] + equil = kwargs["equil"] + else: + use_mu = False + if in_physical: + if not (dim_1 in ["e1", "e2", "e3"] and dim_2 in ["e1", "e2", "e3"]): + AssertionError( + 'To perform a plot in physical space you must use two space axes (dim_1, dim_2 in ["e1","e2","e3"]).' + ) + if plot_3D and not in_physical: + AssertionError("To perform a 3D plot you must plot in physical space (activate in_physical).") + assert 0 <= proj_axis[0] < proj_axis[1] < 3 + if dim_2 is None: + if dim_1 == "e1": + axe_to_plot = 0 + elif dim_1 == "e2": + axe_to_plot = 1 + elif dim_1 == "e3": + axe_to_plot = 2 + elif dim_1 == "v1": + axe_to_plot = 3 + elif dim_1 == "v2": + axe_to_plot = 4 + elif dim_1 == "v3": + axe_to_plot = 5 + else: + AssertionError("dim_1argument must match an exiting dimension") + if axe_to_plot - 3 > self.vdim: + AssertionError("Coordinate " + dim_1 + " does not exist with this background") + linspace_space = xp.array([0.0]) + integrate_linspace_vel = xp.linspace(0.0, v_lim, integrate_resol) + if axe_to_plot < 3: + plot_linspace = xp.linspace(0.0, 1.0, resol) + else: + plot_linspace = xp.linspace(0.0, v_lim, resol) + tabs = [linspace_space for _ in range(3)] + self.vdim * [integrate_linspace_vel] + for i in range(3): + tabs[i][0] = logical_coord[i] + tabs[axe_to_plot] = plot_linspace + etas = xp.meshgrid(*tabs, indexing="ij") + total_density = self(*etas) + if use_mu and axe_to_plot == 4: + B_tab = equil.b_xyz(etas[0], etas[1], etas[2]) + B_norm_tab = xp.sqrt(B_tab[0] ** 2 + B_tab[1] ** 2 + B_tab[2] ** 2) + print(xp.shape(B_norm_tab), xp.shape(etas[4])) + total_density *= B_norm_tab / etas[4] + plot_linspace = plot_linspace**2 + axes_to_integrate = [i for i in range(3 + self.vdim)] + axes_to_integrate.remove(axe_to_plot) + total_density = xp.max(total_density, tuple(axes_to_integrate)) + fig, ax = plt.subplots(1, 1) + ax.plot(plot_linspace, total_density) + ax.set_xlabel(dim_1) + ax.set_ylabel("density") + ax.set_title("background profile") + plt.show(block=True) + else: + if dim_1 == "e1": + axe_to_plot1 = 0 + elif dim_1 == "e2": + axe_to_plot1 = 1 + elif dim_1 == "e3": + axe_to_plot1 = 2 + elif dim_1 == "v1": + axe_to_plot1 = 3 + elif dim_1 == "v2": + axe_to_plot1 = 4 + elif dim_1 == "v3": + axe_to_plot1 = 5 + else: + AssertionError("dim_1 argument must match an exiting dimension") + if dim_2 == "e1": + axe_to_plot2 = 0 + elif dim_2 == "e2": + axe_to_plot2 = 1 + elif dim_2 == "e3": + axe_to_plot2 = 2 + elif dim_2 == "v1": + axe_to_plot2 = 3 + elif dim_2 == "v2": + axe_to_plot2 = 4 + elif dim_2 == "v3": + axe_to_plot2 = 5 + else: + AssertionError("dim_2 argument must match an exiting dimension") + if axe_to_plot2 == axe_to_plot1: + AssertionError("You must specify different dimensions for dim_1 and dim_2") + integrate_linspace_vel = xp.linspace(0.0, v_lim, integrate_resol) + tabs = [xp.array([logical_coord[i]]) for i in range(3)] + self.vdim * [integrate_linspace_vel] + if axe_to_plot1 < 3: + plot_linspace1 = xp.linspace(0.0, 1.0, resol) + elif axe_to_plot1 != 4 or (not use_mu): + plot_linspace1 = xp.linspace(0.0, v_lim, resol) + else: + plot_linspace1 = xp.linspace(0.0, xp.sqrt(v_lim), resol) + if axe_to_plot2 < 3: + plot_linspace2 = xp.linspace(0.0, 1.0, resol) + elif axe_to_plot2 != 4 or (not use_mu): + plot_linspace2 = xp.linspace(0.0, v_lim, resol) + else: + plot_linspace2 = xp.linspace(0.0, xp.sqrt(v_lim), resol) + tabs[axe_to_plot1] = plot_linspace1 + tabs[axe_to_plot2] = plot_linspace2 + etas = xp.meshgrid(*tabs, indexing="ij") + if in_physical: + physical_coords = domain(*tabs[:3]) + physical_coords = list(physical_coords) + for i in range(3): + physical_coords[i] = physical_coords[i][tuple([...] + self.vdim * [None])] + physical_coords[i] = xp.broadcast_to(physical_coords[i], etas[0].shape) + for i in range(self.vdim): + physical_coords.append(etas[i + 3]) + total_density = self(*etas) # domain.push((lambda x, y, z:self(x,y,z,*etas[3:])), *tabs[:3], kind='v') + else: + physical_coords = list(etas) + total_density = self(*etas) + if use_mu: + if axe_to_plot1 == 4 or axe_to_plot2 == 4: + B_tab = equil.b_xyz(etas[0], etas[1], etas[2]) + B_norm_tab = xp.sqrt(B_tab[0] ** 2 + B_tab[1] ** 2 + B_tab[2] ** 2) + print(xp.shape(B_norm_tab), xp.shape(etas[4])) + total_density *= B_norm_tab / etas[4] + physical_coords[4] = physical_coords[4] ** 2 + axes_to_integrate = [i for i in range(3 + self.vdim)] + axes_to_integrate.remove(axe_to_plot1) + axes_to_integrate.remove(axe_to_plot2) + total_density = xp.max(total_density, tuple(axes_to_integrate)) + id_dim = [0] * len(etas) + id_dim[axe_to_plot1] = slice(None) + id_dim[axe_to_plot2] = slice(None) + if not plot_3D: + if in_physical: + X = physical_coords[proj_axis[0]][tuple(id_dim)] + Y = physical_coords[proj_axis[1]][tuple(id_dim)] + else: + X = physical_coords[axe_to_plot1][tuple(id_dim)] + Y = physical_coords[axe_to_plot2][tuple(id_dim)] + fig, ax = plt.subplots() + for_color = ax.pcolor(X, Y, total_density) + if in_physical: + ax.set_xlabel(["x", "y", "z"][proj_axis[0]]) + ax.set_ylabel(["x", "y", "z"][proj_axis[1]]) + else: + if use_mu: + if dim_1 == "v2": + dim_1 = "mu" + if dim_2 == "v2": + dim_2 = "mu" + ax.set_xlabel(dim_1) + ax.set_ylabel(dim_2) + else: + print(xp.shape(physical_coords[0][tuple(id_dim)]), xp.shape(total_density)) + norm = Normalize(total_density.min(), total_density.max() + 0.01) + colors = cm.viridis(norm(total_density)) + fig, ax = plt.subplots(subplot_kw={"projection": "3d"}) + for_color = ax.plot_surface( + X=physical_coords[0][tuple(id_dim)], + Y=physical_coords[1][tuple(id_dim)], + Z=physical_coords[2][tuple(id_dim)], + facecolors=colors, + ) + ax.set_xlabel("x") + ax.set_ylabel("y") + ax.set_zlabel("z") + if title is None: + ax.set_title(f"Density in ({dim_1}, {dim_2}) space") + else: + ax.set_title(title) + fig.colorbar(for_color) + plt.show(block=True) + def __add__(self, other_f0): return SumKineticBackground(self, other_f0) diff --git a/src/struphy/kinetic_background/maxwellians.py b/src/struphy/kinetic_background/maxwellians.py index d9e34dcab..dddf3881a 100644 --- a/src/struphy/kinetic_background/maxwellians.py +++ b/src/struphy/kinetic_background/maxwellians.py @@ -6,7 +6,9 @@ from struphy.fields_background.base import FluidEquilibriumWithB from struphy.fields_background.equils import set_defaults +from struphy.geometry.base import Domain from struphy.initial.base import Perturbation +from struphy.io.options import LiteralOptions from struphy.kinetic_background.base import Maxwellian @@ -300,6 +302,39 @@ def vth(self, eta1, eta2, eta3): out += [self._evaluate_moment(eta1, eta2, eta3, name="vth_perp")] return [ou * mom_fac for ou, mom_fac in zip(out, self.moment_factors["vth"])] + def plot_density_profile( + self, + dim_1: LiteralOptions.DimensionToPlot = "e1", + dim_2: LiteralOptions.DimensionToPlot | None = None, + v_lim: float = 5.0, + resol: int = 100, + integrate_resol: int = 10, + logical_coord: tuple[float] = (0.5, 0.5, 0.5), + in_physical: bool = False, + domain: Domain | None = None, + proj_axis: tuple[float,] = (0, 1), + plot_3D: bool = False, + title: str | None = None, + use_mu: bool = False, + equil: FluidEquilibriumWithB | None = None, + ): + if equil is None: + equil = self.equil + super().plot_density_profile( + dim_1, + dim_2, + v_lim, + resol, + integrate_resol, + logical_coord, + in_physical, + domain, + proj_axis, + plot_3D, + use_mu=use_mu, + equil=equil, + ) + class CanonicalMaxwellian: r"""canonical Maxwellian distribution function. diff --git a/src/struphy/kinetic_background/tests/test_base.py b/src/struphy/kinetic_background/tests/test_base.py index df32b44d2..bdfae8b3f 100644 --- a/src/struphy/kinetic_background/tests/test_base.py +++ b/src/struphy/kinetic_background/tests/test_base.py @@ -84,5 +84,58 @@ def test_kinetic_background_magics(show_plot=False): plt.show() +def test_plotting_function(): + + import cunumpy as xp + + from struphy import domains, equils, maxwellians + + # definition of test functions + l, m, n = 3, 4, 5 + + def n_init(*etas): + if len(etas) == 1: + e1, e2, e3 = etas[0][:, 0], etas[0][:, 1], etas[0][:, 1] + else: + assert len(etas) == 3 + e1, e2, e3 = etas[0], etas[1], etas[2] + return 1 + 0.5 * xp.cos(2 * xp.pi * e1 * l) * xp.cos(2 * xp.pi * e2 * m) * xp.cos(2 * xp.pi * e3 * n) + + def vth(*etas): + if len(etas) == 1: + e1, e2, e3 = etas[0][:, 0], etas[0][:, 1], etas[0][:, 1] + else: + assert len(etas) == 3 + e1, e2, e3 = etas[0], etas[1], etas[2] + return 1 + 0.2 * xp.cos(2 * xp.pi * e1 * l) * xp.cos(2 * xp.pi * e2 * m) * xp.cos(2 * xp.pi * e3 * n) + + # Testing with GyroMaxwellian2D: + equil = equils.HomogenSlab(B0x=0.0, B0y=0.0, B0z=1.0) + background = maxwellians.GyroMaxwellian2D(n=(n_init, None), vth_para=(vth, None), vth_perp=(vth, None), equil=equil) + background.plot_density_profile("e1") + background.plot_density_profile("e2") + background.plot_density_profile("e3") + background.plot_density_profile("v1") + background.plot_density_profile("v2") + background.plot_density_profile("e1", "e2") + background.plot_density_profile("e1", "e2", domain=domains.HollowCylinder(), proj_axis=(0, 1), in_physical=True) + background.plot_density_profile("e1", "e2", domain=domains.HollowTorus(), proj_axis=(1, 2), in_physical=True) + background.plot_density_profile( + "e1", "e2", domain=domains.HollowTorus(), proj_axis=(0, 2), in_physical=True, plot_3D=True + ) + background.plot_density_profile( + "e2", "e3", domain=domains.HollowTorus(), proj_axis=(1, 2), in_physical=True, plot_3D=True + ) + background.plot_density_profile("v1", "v2") + background.plot_density_profile("v1", "v2", use_mu=True) + + # Testing with Maxwellian3D: + background = maxwellians.Maxwellian3D(n=(n_init, None), vth1=(vth, None), vth2=(vth, None), vth3=(vth, None)) + background.plot_density_profile("v1", "v2") + background.plot_density_profile("v1", "v3") + background.plot_density_profile("e1", "v3") + + if __name__ == "__main__": - test_kinetic_background_magics(show_plot=True) + # test_kinetic_background_magics(show_plot=True) + test_plotting_function() From 659e428eaec9c39a1f1edf95660d9781142bd0e0 Mon Sep 17 00:00:00 2001 From: emile grivet Date: Thu, 7 May 2026 16:01:41 +0200 Subject: [PATCH 2/5] corrections in plotting function and removal of output cells in tutorial 3 --- src/struphy/io/options.py | 1 + src/struphy/kinetic_background/base.py | 46 +++++----- src/struphy/kinetic_background/maxwellians.py | 5 +- tutorials/tutorial_03_test_particles.ipynb | 83 ++++++++++++++++++- 4 files changed, 104 insertions(+), 31 deletions(-) diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index 6ecfff7c7..ece2ea02f 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -52,6 +52,7 @@ class LiteralOptions: # fields background BackgroundTypes = Literal["LogicalConst", "FluidEquilibrium"] + KineticDimensionsToPlot = Literal["e1", "e2", "e3", "v1", "v2", "v3"] # models ModelTypes = Literal["Toy", "Kinetic", "Fluid", "Hybrid"] diff --git a/src/struphy/kinetic_background/base.py b/src/struphy/kinetic_background/base.py index 940163eb2..1bc30a014 100644 --- a/src/struphy/kinetic_background/base.py +++ b/src/struphy/kinetic_background/base.py @@ -116,8 +116,8 @@ def __call__(self, *args): def plot_density_profile( self, - dim_1: LiteralOptions.DimensionToPlot = "e1", - dim_2: LiteralOptions.DimensionToPlot | None = None, + dim_1: LiteralOptions.KineticDimensionsToPlot = "e1", + dim_2: LiteralOptions.KineticDimensionsToPlot | None = None, v_lim: float = 5.0, resol: int = 100, integrate_resol: int = 50, @@ -127,14 +127,15 @@ def plot_density_profile( proj_axis: tuple[float,] = (0, 1), plot_3D: bool = False, title: str | None = None, - **kwargs, + use_mu = False, + equil = None, ): """ - Plots the density profile of an initial condition (or a background) from the phase space distribution. The projection can either be 1D or 2D, in the logical space or in cartesian. + Plots the density profile (slices) of the phase space background distribution. The slice can either be 1D or 2D, in logical or in Cartesian coordinates. Parameters ---------- - dim_1, dim_2 : LiteralOptions = ["e1","e2","e3","v1","v2","v3"] + dim_1, dim_2 : LiteralOptions.KineticDimensionsToPlot = ["e1","e2","e3","v1","v2","v3"] The axes used in the projection, they refere to logical space axes. If dim_2 is not defined the projection is 1D, it is 2D if dim_2 is attributed. v_lim : float = 5.0 @@ -150,23 +151,19 @@ def plot_density_profile( Refere to the default coordinate (in logical space) attributed to each axe which is not used in the projection. in_physical : bool = False - Specify if the result is plotted in logical coordinates or in cartesian coordinates, has a effect in 2D plotting. If True, you must specify a domain. + Specify if the result is plotted in logical coordinates or in Cartesian coordinates, has a effect in 2D plotting. If True, you must specify a domain. domain : Domain | None = None Domain used to plot the density if in_physical=True. proj_axis : tuple[float] = (0,1) - Axes of the cartesian coordinates used to plot the density: 0->x, 1->y, 2->z. + Axes of the Cartesian coordinates used to plot the density: 0->x, 1->y, 2->z. I you do not see the density profile in 2D, you may change these axes. plot_3D : bool = False Plots a surface in a 3D environment. Only for physical projection. """ - if "use_mu" in kwargs: - use_mu = kwargs["use_mu"] - equil = kwargs["equil"] - else: - use_mu = False + integrate_resol = min(integrate_resol, 100) # to avoid memory issues if in_physical: if not (dim_1 in ["e1", "e2", "e3"] and dim_2 in ["e1", "e2", "e3"]): AssertionError( @@ -197,22 +194,21 @@ def plot_density_profile( if axe_to_plot < 3: plot_linspace = xp.linspace(0.0, 1.0, resol) else: - plot_linspace = xp.linspace(0.0, v_lim, resol) - tabs = [linspace_space for _ in range(3)] + self.vdim * [integrate_linspace_vel] + plot_linspace = xp.linspace(-v_lim, v_lim, resol) + tabs = 3 * [linspace_space] + self.vdim * [integrate_linspace_vel] for i in range(3): tabs[i][0] = logical_coord[i] tabs[axe_to_plot] = plot_linspace - etas = xp.meshgrid(*tabs, indexing="ij") + etas = xp.abs(xp.meshgrid(*tabs, indexing="ij")) total_density = self(*etas) if use_mu and axe_to_plot == 4: B_tab = equil.b_xyz(etas[0], etas[1], etas[2]) B_norm_tab = xp.sqrt(B_tab[0] ** 2 + B_tab[1] ** 2 + B_tab[2] ** 2) - print(xp.shape(B_norm_tab), xp.shape(etas[4])) total_density *= B_norm_tab / etas[4] plot_linspace = plot_linspace**2 axes_to_integrate = [i for i in range(3 + self.vdim)] axes_to_integrate.remove(axe_to_plot) - total_density = xp.max(total_density, tuple(axes_to_integrate)) + total_density = xp.mean(total_density, tuple(axes_to_integrate)) fig, ax = plt.subplots(1, 1) ax.plot(plot_linspace, total_density) ax.set_xlabel(dim_1) @@ -233,7 +229,7 @@ def plot_density_profile( elif dim_1 == "v3": axe_to_plot1 = 5 else: - AssertionError("dim_1 argument must match an exiting dimension") + AssertionError("dim_1 argument must match an existing dimension") if dim_2 == "e1": axe_to_plot2 = 0 elif dim_2 == "e2": @@ -247,7 +243,7 @@ def plot_density_profile( elif dim_2 == "v3": axe_to_plot2 = 5 else: - AssertionError("dim_2 argument must match an exiting dimension") + AssertionError("dim_2 argument must match an existing dimension") if axe_to_plot2 == axe_to_plot1: AssertionError("You must specify different dimensions for dim_1 and dim_2") integrate_linspace_vel = xp.linspace(0.0, v_lim, integrate_resol) @@ -255,13 +251,13 @@ def plot_density_profile( if axe_to_plot1 < 3: plot_linspace1 = xp.linspace(0.0, 1.0, resol) elif axe_to_plot1 != 4 or (not use_mu): - plot_linspace1 = xp.linspace(0.0, v_lim, resol) + plot_linspace1 = xp.linspace(-v_lim, v_lim, resol) else: plot_linspace1 = xp.linspace(0.0, xp.sqrt(v_lim), resol) if axe_to_plot2 < 3: plot_linspace2 = xp.linspace(0.0, 1.0, resol) elif axe_to_plot2 != 4 or (not use_mu): - plot_linspace2 = xp.linspace(0.0, v_lim, resol) + plot_linspace2 = xp.linspace(-v_lim, v_lim, resol) else: plot_linspace2 = xp.linspace(0.0, xp.sqrt(v_lim), resol) tabs[axe_to_plot1] = plot_linspace1 @@ -275,21 +271,20 @@ def plot_density_profile( physical_coords[i] = xp.broadcast_to(physical_coords[i], etas[0].shape) for i in range(self.vdim): physical_coords.append(etas[i + 3]) - total_density = self(*etas) # domain.push((lambda x, y, z:self(x,y,z,*etas[3:])), *tabs[:3], kind='v') + total_density = self(*xp.abs(etas)) # domain.push((lambda x, y, z:self(x,y,z,*etas[3:])), *tabs[:3], kind='v') else: physical_coords = list(etas) - total_density = self(*etas) + total_density = self(*xp.abs(etas)) if use_mu: if axe_to_plot1 == 4 or axe_to_plot2 == 4: B_tab = equil.b_xyz(etas[0], etas[1], etas[2]) B_norm_tab = xp.sqrt(B_tab[0] ** 2 + B_tab[1] ** 2 + B_tab[2] ** 2) - print(xp.shape(B_norm_tab), xp.shape(etas[4])) total_density *= B_norm_tab / etas[4] physical_coords[4] = physical_coords[4] ** 2 axes_to_integrate = [i for i in range(3 + self.vdim)] axes_to_integrate.remove(axe_to_plot1) axes_to_integrate.remove(axe_to_plot2) - total_density = xp.max(total_density, tuple(axes_to_integrate)) + total_density = xp.mean(total_density, tuple(axes_to_integrate)) id_dim = [0] * len(etas) id_dim[axe_to_plot1] = slice(None) id_dim[axe_to_plot2] = slice(None) @@ -314,7 +309,6 @@ def plot_density_profile( ax.set_xlabel(dim_1) ax.set_ylabel(dim_2) else: - print(xp.shape(physical_coords[0][tuple(id_dim)]), xp.shape(total_density)) norm = Normalize(total_density.min(), total_density.max() + 0.01) colors = cm.viridis(norm(total_density)) fig, ax = plt.subplots(subplot_kw={"projection": "3d"}) diff --git a/src/struphy/kinetic_background/maxwellians.py b/src/struphy/kinetic_background/maxwellians.py index dddf3881a..ed19be5c6 100644 --- a/src/struphy/kinetic_background/maxwellians.py +++ b/src/struphy/kinetic_background/maxwellians.py @@ -304,8 +304,8 @@ def vth(self, eta1, eta2, eta3): def plot_density_profile( self, - dim_1: LiteralOptions.DimensionToPlot = "e1", - dim_2: LiteralOptions.DimensionToPlot | None = None, + dim_1: LiteralOptions.KineticDimensionsToPlot = "e1", + dim_2: LiteralOptions.KineticDimensionsToPlot | None = None, v_lim: float = 5.0, resol: int = 100, integrate_resol: int = 10, @@ -331,6 +331,7 @@ def plot_density_profile( domain, proj_axis, plot_3D, + title, use_mu=use_mu, equil=equil, ) diff --git a/tutorials/tutorial_03_test_particles.ipynb b/tutorials/tutorial_03_test_particles.ipynb index 77c7e0d7c..799b910fc 100644 --- a/tutorials/tutorial_03_test_particles.ipynb +++ b/tutorials/tutorial_03_test_particles.ipynb @@ -359,7 +359,84 @@ "ax.set_aspect(\"equal\")\n", "plt.xlabel(\"x\")\n", "plt.ylabel(\"y\")\n", - "plt.title(\"sim_2: draw uniform on disc\");" + "plt.title(\"sim_2: draw uniform on disc\")" + ] + }, + { + "cell_type": "markdown", + "id": "7bd5ee22", + "metadata": {}, + "source": [ + "Sometimes you need to load particles more uniformly. It is usefull to reach an homogeneous description of the plasma with a limited number of particles.\n", + "\n", + "It is performed using `loading=\"sobol_standard\"` or `loading=\"sobol_antithetic\"` in the `LoadingParameters` function:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f4643154", + "metadata": {}, + "outputs": [], + "source": [ + "env_3 = EnvironmentOptions(sim_folder=\"sim_3\")\n", + "model_3 = Vlasov()\n", + "sim_3 = Simulation.spawn_sister(sim, env=env_3, model=model_3)\n", + "\n", + "# species parameters\n", + "loading_params_3 = LoadingParameters(Np=1000, spatial=\"disc\", loading=\"sobol_standard\")\n", + "\n", + "model_3.kinetic_ions.set_markers(\n", + " loading_params=loading_params_3, weights_params=weights_params, boundary_params=boundary_params\n", + ")\n", + "model_3.kinetic_ions.set_sorting_boxes()\n", + "\n", + "model_3.kinetic_ions.set_save_data(n_markers=1.0)\n", + "\n", + "# propagator options\n", + "model_3.propagators.push_vxb.options = model_3.propagators.push_vxb.Options()\n", + "model_3.propagators.push_eta.options = model_3.propagators.push_eta.Options()\n", + "\n", + "# initial conditions (background + perturbation)\n", + "model_3.kinetic_ions.var.add_background(background)\n", + "\n", + "sim_3.run()\n", + "sim_3.pproc()\n", + "sim_3.load_plotting_data()\n", + "\n", + "fig = plt.figure(figsize=(15, 6))\n", + "\n", + "orbits_standard = sim_3.orbits.kinetic_ions\n", + "\n", + "plt.subplot(1, 3, 1)\n", + "plt.scatter(orbits[0, :, 0], orbits[0, :, 1], s=2.0)\n", + "circle1 = plt.Circle((0, 0), a2, color=\"k\", fill=False)\n", + "ax = plt.gca()\n", + "ax.add_patch(circle1)\n", + "ax.set_aspect(\"equal\")\n", + "plt.xlabel(\"x\")\n", + "plt.ylabel(\"y\")\n", + "plt.title(\"sim_1: draw uniform in logical space\")\n", + "\n", + "plt.subplot(1, 3, 2)\n", + "plt.scatter(orbits_uni[0, :, 0], orbits_uni[0, :, 1], s=2.0)\n", + "circle2 = plt.Circle((0, 0), a2, color=\"k\", fill=False)\n", + "ax = plt.gca()\n", + "ax.add_patch(circle2)\n", + "ax.set_aspect(\"equal\")\n", + "plt.xlabel(\"x\")\n", + "plt.ylabel(\"y\")\n", + "plt.title(\"sim_2: draw uniform on disc\")\n", + "\n", + "plt.subplot(1, 3, 3)\n", + "plt.scatter(orbits_standard[0, :, 0], orbits_standard[0, :, 1], s=2.0)\n", + "circle3 = plt.Circle((0, 0), a2, color=\"k\", fill=False)\n", + "ax = plt.gca()\n", + "ax.add_patch(circle3)\n", + "ax.set_aspect(\"equal\")\n", + "plt.xlabel(\"x\")\n", + "plt.ylabel(\"y\")\n", + "plt.title(\"sim_3: draw using sobol_standard\")" ] }, { @@ -1292,7 +1369,7 @@ ], "metadata": { "kernelspec": { - "display_name": "env (3.12.3)", + "display_name": ".venv (3.13.12)", "language": "python", "name": "python3" }, @@ -1306,7 +1383,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.12.3" + "version": "3.13.12" } }, "nbformat": 4, From afed38e7fc2e4c2188918d31ffafb5e90aa61bdd Mon Sep 17 00:00:00 2001 From: emile grivet Date: Thu, 7 May 2026 16:02:49 +0200 Subject: [PATCH 3/5] corrections in plotting function and removal of output cells in tutorial 3 --- src/struphy/kinetic_background/base.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/struphy/kinetic_background/base.py b/src/struphy/kinetic_background/base.py index 1bc30a014..4188ae0ad 100644 --- a/src/struphy/kinetic_background/base.py +++ b/src/struphy/kinetic_background/base.py @@ -127,8 +127,8 @@ def plot_density_profile( proj_axis: tuple[float,] = (0, 1), plot_3D: bool = False, title: str | None = None, - use_mu = False, - equil = None, + use_mu=False, + equil=None, ): """ Plots the density profile (slices) of the phase space background distribution. The slice can either be 1D or 2D, in logical or in Cartesian coordinates. @@ -271,7 +271,7 @@ def plot_density_profile( physical_coords[i] = xp.broadcast_to(physical_coords[i], etas[0].shape) for i in range(self.vdim): physical_coords.append(etas[i + 3]) - total_density = self(*xp.abs(etas)) # domain.push((lambda x, y, z:self(x,y,z,*etas[3:])), *tabs[:3], kind='v') + total_density = self(*xp.abs(etas)) else: physical_coords = list(etas) total_density = self(*xp.abs(etas)) From 71a4c284f31950c9762479e6de461cabed198c84 Mon Sep 17 00:00:00 2001 From: emile grivet Date: Fri, 8 May 2026 12:58:14 +0200 Subject: [PATCH 4/5] use of equil.absB0 instead of norm(equil.b_xyz) --- src/struphy/kinetic_background/base.py | 36 ++++++++++++++++--- .../kinetic_background/tests/test_base.py | 4 ++- 2 files changed, 35 insertions(+), 5 deletions(-) diff --git a/src/struphy/kinetic_background/base.py b/src/struphy/kinetic_background/base.py index 4188ae0ad..fc09444d1 100644 --- a/src/struphy/kinetic_background/base.py +++ b/src/struphy/kinetic_background/base.py @@ -202,8 +202,22 @@ def plot_density_profile( etas = xp.abs(xp.meshgrid(*tabs, indexing="ij")) total_density = self(*etas) if use_mu and axe_to_plot == 4: - B_tab = equil.b_xyz(etas[0], etas[1], etas[2]) - B_norm_tab = xp.sqrt(B_tab[0] ** 2 + B_tab[1] ** 2 + B_tab[2] ** 2) + B_norm_tab = equil.absB0( + etas[0][tuple([slice(None), slice(None), slice(None)] + self.vdim * [0])], + etas[1][tuple([slice(None), slice(None), slice(None)] + self.vdim * [0])], + etas[2][tuple([slice(None), slice(None), slice(None)] + self.vdim * [0])], + ) + B_norm_tab = xp.broadcast_to( + B_norm_tab[ + tuple( + [ + ..., + ] + + self.vdim * [None] + ) + ], + etas[0].shape, + ) total_density *= B_norm_tab / etas[4] plot_linspace = plot_linspace**2 axes_to_integrate = [i for i in range(3 + self.vdim)] @@ -277,8 +291,22 @@ def plot_density_profile( total_density = self(*xp.abs(etas)) if use_mu: if axe_to_plot1 == 4 or axe_to_plot2 == 4: - B_tab = equil.b_xyz(etas[0], etas[1], etas[2]) - B_norm_tab = xp.sqrt(B_tab[0] ** 2 + B_tab[1] ** 2 + B_tab[2] ** 2) + B_norm_tab = equil.absB0( + etas[0][tuple([slice(None), slice(None), slice(None)] + self.vdim * [0])], + etas[1][tuple([slice(None), slice(None), slice(None)] + self.vdim * [0])], + etas[2][tuple([slice(None), slice(None), slice(None)] + self.vdim * [0])], + ) + B_norm_tab = xp.broadcast_to( + B_norm_tab[ + tuple( + [ + ..., + ] + + self.vdim * [None] + ) + ], + etas[0].shape, + ) total_density *= B_norm_tab / etas[4] physical_coords[4] = physical_coords[4] ** 2 axes_to_integrate = [i for i in range(3 + self.vdim)] diff --git a/src/struphy/kinetic_background/tests/test_base.py b/src/struphy/kinetic_background/tests/test_base.py index bdfae8b3f..340d5740a 100644 --- a/src/struphy/kinetic_background/tests/test_base.py +++ b/src/struphy/kinetic_background/tests/test_base.py @@ -90,6 +90,9 @@ def test_plotting_function(): from struphy import domains, equils, maxwellians + equil = equils.HomogenSlab(B0x=0.0, B0y=0.0, B0z=1.0) + equil.domain = domains.Cuboid() + # definition of test functions l, m, n = 3, 4, 5 @@ -110,7 +113,6 @@ def vth(*etas): return 1 + 0.2 * xp.cos(2 * xp.pi * e1 * l) * xp.cos(2 * xp.pi * e2 * m) * xp.cos(2 * xp.pi * e3 * n) # Testing with GyroMaxwellian2D: - equil = equils.HomogenSlab(B0x=0.0, B0y=0.0, B0z=1.0) background = maxwellians.GyroMaxwellian2D(n=(n_init, None), vth_para=(vth, None), vth_perp=(vth, None), equil=equil) background.plot_density_profile("e1") background.plot_density_profile("e2") From ddfac85f9ed128ee1fb68c1000d171c2a1fed970 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 18 May 2026 11:21:48 +0200 Subject: [PATCH 5/5] trigger ci