From 1c443b2437c7db7715ae3cddf251a1f33c93700d Mon Sep 17 00:00:00 2001 From: MatejTomes Date: Tue, 25 Nov 2025 22:30:51 +0100 Subject: [PATCH 1/5] Add generic distribution function --- cherab/core/distribution.pxd | 8 +++ cherab/core/distribution.pyx | 103 +++++++++++++++++++++++++++++++++++ 2 files changed, 111 insertions(+) diff --git a/cherab/core/distribution.pxd b/cherab/core/distribution.pxd index 062d7a45..af6f3844 100644 --- a/cherab/core/distribution.pxd +++ b/cherab/core/distribution.pxd @@ -19,6 +19,7 @@ from raysect.optical cimport Vector3D from cherab.core.math cimport Function3D, VectorFunction3D +from cherab.core.math.function.float cimport Function6D, autowrap_function6d cdef class DistributionFunction: @@ -45,3 +46,10 @@ cdef class Maxwellian(DistributionFunction): VectorFunction3D _velocity double _atomic_mass + +cdef class GenericDistribution(DistributionFunction): + + cdef readonly: + Function6D _phase_space_density + Function3D _density, _temperature + VectorFunction3D _velocity \ No newline at end of file diff --git a/cherab/core/distribution.pyx b/cherab/core/distribution.pyx index 869cb373..161120a0 100644 --- a/cherab/core/distribution.pyx +++ b/cherab/core/distribution.pyx @@ -25,6 +25,7 @@ from raysect.optical cimport Vector3D cimport cython from cherab.core.math cimport autowrap_function3d, autowrap_vectorfunction3d +from cherab.core.math.function.float cimport autowrap_function6d from cherab.core.utility.constants cimport ELEMENTARY_CHARGE @@ -301,3 +302,105 @@ cdef class Maxwellian(DistributionFunction): return self._density.evaluate(x, y, z) +cdef class GenericDistribution(DistributionFunction): + """ + A generic distribution function. + + This class implements a generic distribution function. The user supplies a 6D function + that provides the phase space density at a given point in 6D phase space, + a 3D function that provides the spatial density, a 3D function that provides the temperature, + and a 3D vector function that provides the bulk velocity. + + ..Warning:: + The consistency of the provided functions is not checked and is the responsibilty of the user. + + :param Function6D phase_space_density: 6D function defining the phase space density in s^3/m^6. + :param Function3D density: 3D function defining the spatial density in m^-3. + :param Function3D temperature: 3D function defining the temperature in eV. + :param VectorFunction3D velocity: 3D vector function defining the bulk velocity in meters per second. + + .. code-block:: pycon + + >>> from cherab.core import GenericDistribution + >>> from cherab.core.math import Function6D, Function3D, VectorFunction3D + >>> + >>> # Setup distribution for a slab of plasma in thermodynamic equilibrium + >>> phase_space_density = Function6D(lambda x, y, z, vx, vy, vz: 1E17 * exp(-(vx**2 + vy**2 + vz**2) / (2 * 1))) + >>> density = Function3D(lambda x, y, z: 1E17) + >>> temperature = Function3D(lambda x, y, z: 1) + >>> velocity = VectorFunction3D(lambda x, y, z: Vector3D(0, 0, 0)) + >>> d0_distribution = GenericDistribution(phase_space_density, density, temperature, velocity) + """ + + def __init__(self, object phase_space_density, object density, object temperature, object velocity): + + super().__init__() + self._phase_space_density = autowrap_function6d(phase_space_density) + self._density = autowrap_function3d(density) + self._temperature = autowrap_function3d(temperature) + self._velocity = autowrap_vectorfunction3d(velocity) + + @cython.cdivision(True) + cdef double evaluate(self, double x, double y, double z, double vx, double vy, double vz) except? -1e999: + """ + Evaluates the phase space density at the specified point in 6D phase space. + + :param x: position in meters + :param y: position in meters + :param z: position in meters + :param vx: velocity in meters per second + :param vy: velocity in meters per second + :param vz: velocity in meters per second + :return: phase space density in s^3/m^6 + """ + + return self._phase_space_density.evaluate(x, y, z, vx, vy, vz) + + cpdef Vector3D bulk_velocity(self, double x, double y, double z): + """ + Evaluates the species' bulk velocity at the specified 3D coordinate. + + :param x: position in meters + :param y: position in meters + :param z: position in meters + :return: velocity vector in m/s + + .. code-block:: pycon + + >>> d0_distribution.bulk_velocity(1, 0, 0) + Vector3D(0.0, 0.0, 0.0) + """ + + return self._velocity.evaluate(x, y, z) + + cpdef double effective_temperature(self, double x, double y, double z) except? -1e999: + """ + + :param x: position in meters + :param y: position in meters + :param z: position in meters + :return: temperature in eV + + .. code-block:: pycon + + >>> d0_distribution.effective_temperature(1, 0, 0) + 1.0 + """ + + return self._temperature.evaluate(x, y, z) + + cpdef double density(self, double x, double y, double z) except? -1e999: + """ + + :param x: position in meters + :param y: position in meters + :param z: position in meters + :return: density in m^-3 + + .. code-block:: pycon + + >>> d0_distribution.density(1, 0, 0) + 1e+17 + """ + + return self._density.evaluate(x, y, z) \ No newline at end of file From bd869cc04501bead0efe780deaeea702e0b13f72 Mon Sep 17 00:00:00 2001 From: MatejTomes Date: Tue, 25 Nov 2025 22:31:46 +0100 Subject: [PATCH 2/5] Add tests to GenericDistribution, ZeroDistribution move Maxwellian tests from test_maxwellian.py to test_distribution.py --- cherab/core/tests/test_distribution.py | 325 +++++++++++++++++++++++++ cherab/core/tests/test_maxwellian.py | 109 --------- 2 files changed, 325 insertions(+), 109 deletions(-) create mode 100644 cherab/core/tests/test_distribution.py delete mode 100644 cherab/core/tests/test_maxwellian.py diff --git a/cherab/core/tests/test_distribution.py b/cherab/core/tests/test_distribution.py new file mode 100644 index 00000000..caa1cfd0 --- /dev/null +++ b/cherab/core/tests/test_distribution.py @@ -0,0 +1,325 @@ +# Copyright 2016-2018 Euratom +# Copyright 2016-2018 United Kingdom Atomic Energy Authority +# Copyright 2016-2018 Centro de Investigaciones Energéticas, Medioambientales y Tecnológicas +# +# Licensed under the EUPL, Version 1.1 or – as soon they will be approved by the +# European Commission - subsequent versions of the EUPL (the "Licence"); +# You may not use this work except in compliance with the Licence. +# You may obtain a copy of the Licence at: +# +# https://joinup.ec.europa.eu/software/page/eupl5 +# +# Unless required by applicable law or agreed to in writing, software distributed +# under the Licence is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR +# CONDITIONS OF ANY KIND, either express or implied. +# +# See the Licence for the specific language governing permissions and limitations +# under the Licence. + +import unittest + +import numpy as np + +from cherab.core.distribution import ZeroDistribution, GenericDistribution, Maxwellian +from raysect.core import Vector3D + +ATOMIC_MASS = 1.66053906660e-27 +ELEMENTARY_CHARGE = 1.602176634e-19 + + +# Note: DistributionFunction is a cdef class (abstract base class) that cannot be +# directly instantiated or subclassed from Python. The abstract methods raise +# NotImplementedError, which is tested implicitly through the concrete implementations +# (ZeroDistribution, Maxwellian, GenericDistribution) that inherit from it. + + +class TestZeroDistribution(unittest.TestCase): + """ + Test cases for the ZeroDistribution class. + + ZeroDistribution should return zero for all distribution properties. + """ + + def setUp(self): + self.distribution = ZeroDistribution() + self.x = np.linspace(-10, 10, 5) # m + self.y = np.linspace(-10, 10, 5) # m + self.z = np.linspace(-10, 10, 5) # m + self.vx = np.linspace(-10e5, 10e5, 5) # m/s + self.vy = np.linspace(-10e5, 10e5, 5) # m/s + self.vz = np.linspace(-10e5, 10e5, 5) # m/s + + def tearDown(self): + pass + + def test_call_returns_zero(self): + """Test that __call__() returns zero for all inputs.""" + for x in self.x[::2]: # Test subset to avoid long execution time + for y in self.y[::2]: + for z in self.z[::2]: + for vx in self.vx[::2]: + for vy in self.vy[::2]: + for vz in self.vz[::2]: + result = self.distribution(x, y, z, vx, vy, vz) + self.assertEqual(result, 0.0, + msg='__call__() should return 0.0 at ({}, {}, {}, {}, {}, {}).'.format( + x, y, z, vx, vy, vz)) + + def test_bulk_velocity_returns_zero_vector(self): + """Test that bulk_velocity() returns zero vector for all positions.""" + for x in self.x: + for y in self.y: + for z in self.z: + velocity = self.distribution.bulk_velocity(x, y, z) + self.assertAlmostEqual(velocity.x, 0.0, delta=1e-10, + msg='bulk_velocity().x should be 0.0 at ({}, {}, {}).'.format(x, y, z)) + self.assertAlmostEqual(velocity.y, 0.0, delta=1e-10, + msg='bulk_velocity().y should be 0.0 at ({}, {}, {}).'.format(x, y, z)) + self.assertAlmostEqual(velocity.z, 0.0, delta=1e-10, + msg='bulk_velocity().z should be 0.0 at ({}, {}, {}).'.format(x, y, z)) + + def test_effective_temperature_returns_zero(self): + """Test that effective_temperature() returns zero for all positions.""" + for x in self.x: + for y in self.y: + for z in self.z: + temperature = self.distribution.effective_temperature(x, y, z) + self.assertEqual(temperature, 0.0, + msg='effective_temperature() should return 0.0 at ({}, {}, {}).'.format(x, y, z)) + + def test_density_returns_zero(self): + """Test that density() returns zero for all positions.""" + for x in self.x: + for y in self.y: + for z in self.z: + density = self.distribution.density(x, y, z) + self.assertEqual(density, 0.0, + msg='density() should return 0.0 at ({}, {}, {}).'.format(x, y, z)) + + +class TestGenericDistribution(unittest.TestCase): + """ + Test cases for the GenericDistribution class. + + GenericDistribution allows users to provide custom 6D phase space density, + density, temperature, and velocity functions. + """ + + def setUp(self): + self.x = np.linspace(-5, 5, 3) # m + self.y = np.linspace(-5, 5, 3) # m + self.z = np.linspace(-5, 5, 3) # m + self.vx = np.linspace(-5e5, 5e5, 3) # m/s + self.vy = np.linspace(-5e5, 5e5, 3) # m/s + self.vz = np.linspace(-5e5, 5e5, 3) # m/s + + # Define atomic mass for Gaussian distribution (using deuterium mass) + self.atomic_mass = 2 * ATOMIC_MASS # kg + + # Define shared density and temperature functions for 3D Gaussian distribution + self.density = lambda x, y, z: 1e20 * (1 + 0.1 * np.sin(x) * np.sin(y) * np.sin(z)) # m^-3 + self.temperature = lambda x, y, z: 1e3 * (1 + 0.1 * np.sin(x + 1) * np.sin(y + 1) * np.sin(z + 1)) # eV + self.velocity = lambda x, y, z: Vector3D(1e5 * x, 2e5 * y, 3e5 * z) # m/s + + + # Define 3D Gaussian phase space density function using density and temperature + # This implements a Maxwellian distribution: f = n * (m/(2*pi*e*T))^(3/2) * exp(-m*v^2/(2*e*T)) + def phase_space_density_gaussian(x, y, z, vx, vy, vz): + n = self.density(x, y, z) + T = self.temperature(x, y, z) + m = self.atomic_mass + + # Thermal velocity spread squared + sigma_sq = T * ELEMENTARY_CHARGE / m # (m/s)^2 + + # Velocity magnitude squared (assuming zero bulk velocity for simplicity) + v_sq = vx**2 + vy**2 + vz**2 + + # Normalization factor + norm = (m / (2 * np.pi * ELEMENTARY_CHARGE * T)) ** 1.5 + + # Gaussian distribution + return n * norm * np.exp(-v_sq / (2 * sigma_sq)) + + self.phase_space_density = phase_space_density_gaussian + + def test_bulk_velocity(self): + """Test that bulk_velocity() returns the correct velocity vector.""" + # Define velocity function + + distribution = GenericDistribution(self.phase_space_density, self.density, self.temperature, self.velocity) + + for x in self.x: + for y in self.y: + for z in self.z: + result = distribution.bulk_velocity(x, y, z) + expected = self.velocity(x, y, z) + self.assertAlmostEqual(result.x, expected.x, delta=1e-10, + msg='bulk_velocity().x is wrong at ({}, {}, {}).'.format(x, y, z)) + self.assertAlmostEqual(result.y, expected.y, delta=1e-10, + msg='bulk_velocity().y is wrong at ({}, {}, {}).'.format(x, y, z)) + self.assertAlmostEqual(result.z, expected.z, delta=1e-10, + msg='bulk_velocity().z is wrong at ({}, {}, {}).'.format(x, y, z)) + + def test_effective_temperature(self): + """Test that effective_temperature() returns the correct temperature.""" + + distribution = GenericDistribution(self.phase_space_density, self.density, self.temperature, self.velocity) + + for x in self.x: + for y in self.y: + for z in self.z: + result = distribution.effective_temperature(x, y, z) + expected = self.temperature(x, y, z) + self.assertAlmostEqual(result, expected, delta=1e-10, + msg='effective_temperature() is wrong at ({}, {}, {}).'.format(x, y, z)) + + def test_density(self): + """Test that density() returns the correct density.""" + + distribution = GenericDistribution(self.phase_space_density, self.density, self.temperature, self.velocity) + + for x in self.x: + for y in self.y: + for z in self.z: + result = distribution.density(x, y, z) + expected = self.density(x, y, z) + self.assertAlmostEqual(result, expected, delta=1e-10, + msg='density() is wrong at ({}, {}, {}).'.format(x, y, z)) + + def test_call(self): + """Test that __call__() returns the correct phase space density using 3D Gaussian.""" + + distribution = GenericDistribution(self.phase_space_density, self.density, self.temperature, self.velocity) + + # Test subset to avoid long execution time + for x in self.x: + for y in self.y: + for z in self.z[::2]: + for vx in self.vx[::2]: + for vy in self.vy[::2]: + for vz in self.vz[::2]: + result = distribution(x, y, z, vx, vy, vz) + expected = self.phase_space_density(x, y, z, vx, vy, vz) + self.assertAlmostEqual(result, expected, delta=1e-10, + msg='__call__() is wrong at ({}, {}, {}, {}, {}, {}).'.format( + x, y, z, vx, vy, vz)) + + def test_float_inputs(self): + """Test GenericDistribution with float inputs instead of lambda functions.""" + # Pass floats directly - they should be converted to constant functions + density_float = 1e20 # m^-3 + temperature_float = 1e3 # eV + velocity_constant = Vector3D(1e5, 2e5, 3e5) # m/s (constant vector function) + phase_space_density_float = 1e17 # s^3/m^6 + + distribution = GenericDistribution(phase_space_density_float, density_float, temperature_float, velocity_constant) + + # Test that all methods work with constant float inputs + for x in self.x: + for y in self.y: + for z in self.z: + # Density should be constant + self.assertAlmostEqual(distribution.density(x, y, z), density_float, delta=1e-10, + msg='density() should return constant value at ({}, {}, {}).'.format(x, y, z)) + # Temperature should be constant + self.assertAlmostEqual(distribution.effective_temperature(x, y, z), temperature_float, delta=1e-10, + msg='effective_temperature() should return constant value at ({}, {}, {}).'.format(x, y, z)) + # Velocity should be constant + vel = distribution.bulk_velocity(x, y, z) + self.assertAlmostEqual(vel.x, velocity_constant.x, delta=1e-10, + msg='bulk_velocity().x should return constant value at ({}, {}, {}).'.format(x, y, z)) + self.assertAlmostEqual(vel.y, velocity_constant.y, delta=1e-10, + msg='bulk_velocity().y should return constant value at ({}, {}, {}).'.format(x, y, z)) + self.assertAlmostEqual(vel.z, velocity_constant.z, delta=1e-10, + msg='bulk_velocity().z should return constant value at ({}, {}, {}).'.format(x, y, z)) + + # Test phase space density is constant + for x in self.x[::2]: + for y in self.y[::2]: + for z in self.z[::2]: + for vx in self.vx[::2]: + for vy in self.vy[::2]: + for vz in self.vz[::2]: + result = distribution(x, y, z, vx, vy, vz) + self.assertAlmostEqual(result, phase_space_density_float, delta=1e-10, + msg='__call__() should return constant value at ({}, {}, {}, {}, {}, {}).'.format( + x, y, z, vx, vy, vz)) + + +class TestMaxwellian(unittest.TestCase): + """ + Test cases for the Maxwellian class. + + Maxwellian implements a Maxwell-Boltzmann distribution function. + """ + + def setUp(self): + self.x = np.linspace(-10, 10, 5) # m + self.y = np.linspace(-10, 10, 5) # m + self.z = np.linspace(-10, 10, 5) # m + self.vx = np.linspace(-10e5, 10e5, 5) # m/s + self.vy = np.linspace(-10e5, 10e5, 5) # m/s + self.vz = np.linspace(-10e5, 10e5, 5) # m/s + + # Define shared density, temperature, velocity, and mass for all tests + self.density = lambda x, y, z: 6e19 * (1 + 0.1 * np.sin(x) * np.sin(y) * np.sin(z)) # m^-3 + self.temperature = lambda x, y, z: 3e3 * (1 + 0.1 * np.sin(x + 1) * np.sin(y + 1) * np.sin(z + 1)) # eV + self.velocity = lambda x, y, z: 1.6e5 * (1 + 0.1 * np.sin(x + 2) * np.sin(y + 2) * np.sin(z + 2)) * Vector3D(1, 2, 3).normalise() # m/s + self.mass = 4 * ATOMIC_MASS # kg + + # Define sigma and phase_space_density for test_value + self.sigma = lambda x, y, z: np.sqrt(self.temperature(x, y, z) * ELEMENTARY_CHARGE / self.mass) # m/s + self.phase_space_density = lambda x, y, z, vx, vy, vz: self.density(x, y, z) / (np.sqrt(2 * np.pi) * self.sigma(x, y, z)) ** 3 \ + * np.exp(-(Vector3D(vx, vy, vz) - self.velocity(x, y, z)).length ** 2 / (2 * self.sigma(x, y, z) ** 2)) # s^3/m^6 + + def tearDown(self): + pass + + def test_bulk_velocity(self): + """Test that bulk_velocity() returns the correct velocity vector.""" + maxwellian = Maxwellian(self.density, self.temperature, self.velocity, self.mass) + + for x in self.x: + for y in self.y: + for z in self.z: + self.assertAlmostEqual(maxwellian.bulk_velocity(x, y, z).x, self.velocity(x, y, z).x, delta=1e-10, + msg='bulk_velocity method gives a wrong value at ({}, {}, {}).'.format(x, y, z)) + self.assertAlmostEqual(maxwellian.bulk_velocity(x, y, z).y, self.velocity(x, y, z).y, delta=1e-10, + msg='bulk_velocity method gives a wrong value at ({}, {}, {}).'.format(x, y, z)) + self.assertAlmostEqual(maxwellian.bulk_velocity(x, y, z).z, self.velocity(x, y, z).z, delta=1e-10, + msg='bulk_velocity method gives a wrong value at ({}, {}, {}).'.format(x, y, z)) + + def test_effective_temperature(self): + """Test that effective_temperature() returns the correct temperature.""" + maxwellian = Maxwellian(self.density, self.temperature, self.velocity, self.mass) + + for x in self.x: + for y in self.y: + for z in self.z: + self.assertAlmostEqual(maxwellian.effective_temperature(x, y, z), self.temperature(x, y, z), delta=1e-10, + msg='effective_temperature method gives a wrong value at ({}, {}, {}).'.format(x, y, z)) + + def test_density(self): + """Test that density() returns the correct density.""" + maxwellian = Maxwellian(self.density, self.temperature, self.velocity, self.mass) + + for x in self.x: + for y in self.y: + for z in self.z: + self.assertAlmostEqual(maxwellian.density(x, y, z), self.density(x, y, z), delta=1e-10, + msg='density method gives a wrong value at ({}, {}, {}).'.format(x, y, z)) + + def test_value(self): + """Test that __call__() returns the correct phase space density.""" + maxwellian = Maxwellian(self.density, self.temperature, self.velocity, self.mass) + + # testing only half the values to avoid huge execution time + for x in self.x[::2]: + for y in self.y[::2]: + for z in self.z[::2]: + for vx in self.vx[::2]: + for vy in self.vy[::2]: + for vz in self.vz[::2]: + self.assertAlmostEqual(maxwellian(x, y, z, vx, vy, vz), self.phase_space_density(x, y, z, vx, vy, vz), delta=1e-10, + msg='call method gives a wrong phase space density at ({}, {}, {}, {}, {}, {}).'.format(x, y, z, vx, vy, vz)) \ No newline at end of file diff --git a/cherab/core/tests/test_maxwellian.py b/cherab/core/tests/test_maxwellian.py deleted file mode 100644 index 67339b4b..00000000 --- a/cherab/core/tests/test_maxwellian.py +++ /dev/null @@ -1,109 +0,0 @@ -# Copyright 2016-2018 Euratom -# Copyright 2016-2018 United Kingdom Atomic Energy Authority -# Copyright 2016-2018 Centro de Investigaciones Energéticas, Medioambientales y Tecnológicas -# -# Licensed under the EUPL, Version 1.1 or – as soon they will be approved by the -# European Commission - subsequent versions of the EUPL (the "Licence"); -# You may not use this work except in compliance with the Licence. -# You may obtain a copy of the Licence at: -# -# https://joinup.ec.europa.eu/software/page/eupl5 -# -# Unless required by applicable law or agreed to in writing, software distributed -# under the Licence is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR -# CONDITIONS OF ANY KIND, either express or implied. -# -# See the Licence for the specific language governing permissions and limitations -# under the Licence. - -import unittest - -import numpy as np - -from cherab.core.distribution import Maxwellian -from raysect.core import Vector3D - -ATOMIC_MASS = 1.66053906660e-27 -ELEMENTARY_CHARGE = 1.602176634e-19 - - -class TestMaxwellian(unittest.TestCase): - - def setUp(self): - self.x = np.linspace(-10, 10, 5) # m - self.y = np.linspace(-10, 10, 5) # m - self.z = np.linspace(-10, 10, 5) # m - self.vx = np.linspace(-10e5, 10e5, 5) # m/s - self.vy = np.linspace(-10e5, 10e5, 5) # m/s - self.vz = np.linspace(-10e5, 10e5, 5) # m/s - - def tearDown(self): - pass - - def test_bulk_velocity(self): - density = lambda x, y, z: 6e19 * (1 + 0.1 * np.sin(x) * np.sin(y) * np.sin(z)) # m^-3 - temperature = lambda x, y, z: 3e3 * (1 + 0.1 * np.sin(x + 1) * np.sin(y + 1) * np.sin(z + 1)) # eV - velocity = lambda x, y, z: 1.6e5 * (1 + 0.1 * np.sin(x + 2) * np.sin(y + 2) * np.sin(z + 2)) * Vector3D(1, 2, 3).normalise() # m/s - mass = 4 * ATOMIC_MASS # kg - maxwellian = Maxwellian(density, temperature, velocity, mass) - - for x in self.x: - for y in self.y: - for z in self.z: - self.assertAlmostEqual(maxwellian.bulk_velocity(x, y, z).x, velocity(x, y, z).x, delta=1e-10, - msg='bulk_velocity method gives a wrong value at ({}, {}, {}).'.format(x, y, z)) - self.assertAlmostEqual(maxwellian.bulk_velocity(x, y, z).y, velocity(x, y, z).y, delta=1e-10, - msg='bulk_velocity method gives a wrong value at ({}, {}, {}).'.format(x, y, z)) - self.assertAlmostEqual(maxwellian.bulk_velocity(x, y, z).z, velocity(x, y, z).z, delta=1e-10, - msg='bulk_velocity method gives a wrong value at ({}, {}, {}).'.format(x, y, z)) - - def test_effective_temperature(self): - density = lambda x, y, z: 6e19 * (1 + 0.1 * np.sin(x) * np.sin(y) * np.sin(z)) # m^-3 - temperature = lambda x, y, z: 3e3 * (1 + 0.1 * np.sin(x + 1) * np.sin(y + 1) * np.sin(z + 1)) # eV - velocity = lambda x, y, z: 1.6e5 * (1 + 0.1 * np.sin(x + 2) * np.sin(y + 2) * np.sin(z + 2)) * Vector3D(1, 2, 3).normalise() # m/s - mass = 4 * ATOMIC_MASS # kg - maxwellian = Maxwellian(density, temperature, velocity, mass) - - for x in self.x: - for y in self.y: - for z in self.z: - self.assertAlmostEqual(maxwellian.effective_temperature(x, y, z), temperature(x, y, z), delta=1e-10, - msg='effective_temperature method gives a wrong value at ({}, {}, {}).'.format(x, y, z)) - - def test_density(self): - density = lambda x, y, z: 6e19 * (1 + 0.1 * np.sin(x) * np.sin(y) * np.sin(z)) # m^-3 - temperature = lambda x, y, z: 3e3 * (1 + 0.1 * np.sin(x + 1) * np.sin(y + 1) * np.sin(z + 1)) # eV - velocity = lambda x, y, z: 1.6e5 * (1 + 0.1 * np.sin(x + 2) * np.sin(y + 2) * np.sin(z + 2)) * Vector3D(1, 2, 3).normalise() # m/s - mass = 4 * ATOMIC_MASS # kg - maxwellian = Maxwellian(density, temperature, velocity, mass) - - for x in self.x: - for y in self.y: - for z in self.z: - self.assertAlmostEqual(maxwellian.density(x, y, z), density(x, y, z), delta=1e-10, - msg='density method gives a wrong value at ({}, {}, {}).'.format(x, y, z)) - - def test_value(self): - density = lambda x, y, z: 6e19 * (1 + 0.1 * np.sin(x) * np.sin(y) * np.sin(z)) # m^-3 - temperature = lambda x, y, z: 3e3 * (1 + 0.1 * np.sin(x + 1) * np.sin(y + 1) * np.sin(z + 1)) # eV - velocity = lambda x, y, z: 1.6e5 * (1 + 0.1 * np.sin(x + 2) * np.sin(y + 2) * np.sin(z + 2)) * Vector3D(1, 2, 3).normalise() # m/s - mass = 4 * ATOMIC_MASS # kg - maxwellian = Maxwellian(density, temperature, velocity, mass) - - sigma = lambda x, y, z: np.sqrt(temperature(x, y, z) * ELEMENTARY_CHARGE / mass) # m/s - phase_space_density = lambda x, y, z, vx, vy, vz: density(x, y, z) / (np.sqrt(2 * np.pi) * sigma(x, y, z)) ** 3 \ - * np.exp(-(Vector3D(vx, vy, vz) - velocity(x, y, z)).length ** 2 / (2 * sigma(x, y, z) ** 2)) # s^3/m^6 - - # testing only half the values to avoid huge execution time - for x in self.x[::2]: - for y in self.y[::2]: - for z in self.z[::2]: - for vx in self.vx[::2]: - for vy in self.vy[::2]: - for vz in self.vz[::2]: - self.assertAlmostEqual(maxwellian(x, y, z, vx, vy, vz), phase_space_density(x, y, z, vx, vy, vz), delta=1e-10, - msg='call method gives a wrong phase space density at ({}, {}, {}, {}, {}, {}).'.format(x, y, z, vx, vy, vz)) - - -if __name__ == '__main__': - unittest.main() \ No newline at end of file From 50cb7cb9f9da757f0267ec705fb0fb247aa71c7d Mon Sep 17 00:00:00 2001 From: MatejTomes Date: Tue, 25 Nov 2025 22:33:38 +0100 Subject: [PATCH 3/5] Add ZeroDistribution, GenericDistribution to docs --- docs/source/plasmas/core_plasma_classes.rst | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/docs/source/plasmas/core_plasma_classes.rst b/docs/source/plasmas/core_plasma_classes.rst index a45a141d..46cd7e64 100644 --- a/docs/source/plasmas/core_plasma_classes.rst +++ b/docs/source/plasmas/core_plasma_classes.rst @@ -30,3 +30,13 @@ Distribution functions :special-members: __call__ :show-inheritance: +.. autoclass:: cherab.core.distribution.GenericDistribution + :members: + :special-members: __call__ + :show-inheritance: + +.. autoclass:: cherab.core.distribution.ZeroDistribution + :members: + :special-members: __call__ + :show-inheritance: + From c901f8aba9422c82073b48ee43aae729f7f1c44e Mon Sep 17 00:00:00 2001 From: MatejTomes Date: Tue, 25 Nov 2025 22:34:34 +0100 Subject: [PATCH 4/5] Update changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6390e90e..11ad74f7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ Release 1.6.0 (TBD) ------------------- API changes: +* Add generic distribution function. (#481) * Add emission model attribute access to line and lineshape . (#294) New: From 1629e68f1df7d646116b6f95f530da38931f6108 Mon Sep 17 00:00:00 2001 From: MatejTomes Date: Wed, 26 Nov 2025 08:54:39 +0100 Subject: [PATCH 5/5] Add demo for the generic distribution --- .../generic_distribution.py | 297 ++++++++++++++++++ 1 file changed, 297 insertions(+) create mode 100644 demos/particle_distribution/generic_distribution.py diff --git a/demos/particle_distribution/generic_distribution.py b/demos/particle_distribution/generic_distribution.py new file mode 100644 index 00000000..cb93576b --- /dev/null +++ b/demos/particle_distribution/generic_distribution.py @@ -0,0 +1,297 @@ +from scipy.constants import atomic_mass, elementary_charge, pi + +import numpy as np + +import matplotlib.pyplot as plt + +from raysect.core.math.function.float import Arg3D, Exp3D, Sqrt3D +from raysect.core.math.function.vector3d import FloatToVector3DFunction3D + +from cherab.core.atomic import deuterium +from cherab.core.distribution import GenericDistribution +from cherab.core.math.function.float import Arg6D, Exp6D, Sqrt6D + + +# To set up a generic distribution, we need to define the following: +# - define 3D scalar function defining the spatial distribution of the effective temperature +# - define 3D vector function defining the spatial distribution of the bulk velocity +# - define 3D scalar function defining the spatial distribution of the density +# - define 6D scalar function defining the phase space density + +# This example creates a toroidally symmetric distribution in R-Z coordinates +# where R = sqrt(X^2 + Y^2) and Z is the vertical coordinate. +# The distribution peaks at R=2, Z=0 with Gaussian-like profiles. + +# initialise the spatial arguments for the 3D functions +x3d, y3d, z3d = Arg3D("x"), Arg3D("y"), Arg3D("z") + +# Calculate R = sqrt(X^2 + Y^2) for 3D functions +r_3d = Sqrt3D(x3d**2 + y3d**2) + +# Peak location in R-Z space +r_peak = 2.0 # meters +z_peak = 0.0 # meters + +# set the properties of the temperature profile +maximum_temperature = 1000 # eV +temperature_peak_width_R = 0.5 # meters +temperature_peak_width_Z = 0.5 # meters + +# set up a 3D gaussian-like temperature profile in R-Z +# The temperature is defined only as a 3D function, +# and is not used within the phase space density function. +# Consistency with the phase space density function is not checked +# and is the responsibility of the user, if required. +temperature_3d = maximum_temperature * Exp3D( + -0.5 + * ( + ((r_3d - r_peak) ** 2 / temperature_peak_width_R**2) + + ((z3d - z_peak) ** 2 / temperature_peak_width_Z**2) + ) +) + +# set the properties of the density profile +maximum_density = 5e19 # m^-3 +density_peak_width_r = 0.5 # meters +density_peak_width_z = 0.5 # meters + +# set up the 3D function defining the spatial density in R-Z +# The density is defined only as a 3D function, +# and is not used within the phase space density function. +# Consistency with the phase space density function is not checked +# and is the responsibility of the user, if required. +density_3d = maximum_density * Exp3D( + -0.5 + * ( + ((r_3d - r_peak) ** 2 / density_peak_width_r**2) + + ((z3d - z_peak) ** 2 / density_peak_width_z**2) + ) +) + +# set the properties of the toroidal rotation velocity profile +# The bulk velocity is defined only as a 3D vector function, +# and is not used within the phase space density function. +# Consistency with the phase space density function is not checked +# and is the responsibility of the user, if required. +maximum_toroidal_velocity = 1e5 # m/s +toroidal_velocity_peak_width_R = 0.5 # meters +toroidal_velocity_peak_width_Z = 0.5 # meters + +# Toroidal velocity profile in R-Z (Gaussian-like) +# The toroidal direction is perpendicular to R and Z +# In Cartesian: v_toroidal * (-y/R, x/R, 0) +v_toroidal_profile_3d = maximum_toroidal_velocity * Exp3D( + -0.5 + * ( + ((r_3d - r_peak) ** 2 / toroidal_velocity_peak_width_R**2) + + ((z3d - z_peak) ** 2 / toroidal_velocity_peak_width_Z**2) + ) +) + +# Convert toroidal velocity to Cartesian components +# vx = -v_toroidal * y/R, vy = v_toroidal * x/R, vz = 0 +# Note: We need to handle the case where R=0, but for this example we assume R>0 +vx_profile = -v_toroidal_profile_3d * y3d / (r_3d + 1e-10) # small epsilon to avoid division by zero +vy_profile = v_toroidal_profile_3d * x3d / (r_3d + 1e-10) +vz_profile = 0.0 +bulk_velocity_profile = FloatToVector3DFunction3D(vx_profile, vy_profile, vz_profile) + +# initialise the arguments for the 6D function +x6d, y6d, z6d, vx6d, vy6d, vz6d = ( + Arg6D("x"), + Arg6D("y"), + Arg6D("z"), + Arg6D("u"), + Arg6D("w"), + Arg6D("v"), +) + +# Calculate R = sqrt(X^2 + Y^2) for 6D functions +r_6d = Sqrt6D(x6d**2 + y6d**2) + +# set the missing parameters of the distribution function +deuterium_mass = deuterium.atomic_weight * atomic_mass + +# re-define the spatial temperature profile with the 6D function parameters +te_6d = maximum_temperature * Exp6D( + -0.5 + * ( + ((r_6d - r_peak) ** 2 / temperature_peak_width_R**2) + + ((z6d - z_peak) ** 2 / temperature_peak_width_Z**2) + ) +) + +# re-define the spatial density profile with the 6D function parameters +density_6d = maximum_density * Exp6D( + -0.5 + * ( + ((r_6d - r_peak) ** 2 / density_peak_width_r**2) + + ((z6d - z_peak) ** 2 / density_peak_width_z**2) + ) +) + +# Toroidal velocity profile redefined with the 6D function parameters +v_toroidal_mean_6d = maximum_toroidal_velocity * Exp6D( + -0.5 + * ( + ((r_6d - r_peak) ** 2 / toroidal_velocity_peak_width_R**2) + + ((z6d - z_peak) ** 2 / toroidal_velocity_peak_width_Z**2) + ) +) * -1.0 + +# Convert toroidal velocity to Cartesian velocity components for the 6D function +# vx_mean = -v_toroidal * y/R, vy_mean = v_toroidal * x/R, vz_mean = 0 +vx_mean_6d = -v_toroidal_mean_6d * y6d / (r_6d + 1e-10) +vy_mean_6d = v_toroidal_mean_6d * x6d / (r_6d + 1e-10) +vz_mean_6d = 0.0 + +# define the 6D distribution function for the bulk population +factor_6d = (deuterium_mass / (2 * pi * elementary_charge * te_6d)) ** 1.5 + +thermal_exponential_6d = Exp6D( + -0.5 + * deuterium_mass + * ((vx6d - vx_mean_6d) ** 2 + (vy6d - vy_mean_6d) ** 2 + (vz6d - vz_mean_6d) ** 2) + / (elementary_charge * te_6d) +) +bulk_pdf_6d = ( + density_6d * factor_6d * thermal_exponential_6d +) # bulk particle distribution function + +# add a population of supra-thermal particles with higher toroidal rotation +# the 3D bulk velocity and temperature functions ignore this population, +# and is the responsibility of the user, if required. +supra_thermal_population_ratio = 0.005 # 10% of the particles are supra-thermal +suprathermal_toroidal_velocity_factor = 20.0 # Supra-thermal particles have 2x toroidal velocity +suprathermal_temperature = 100 # eV (higher temperature for supra-thermal particles) + +# Supra-thermal toroidal velocity profile (higher than bulk) +v_toroidal_suprathermal_6d = ( + maximum_toroidal_velocity + * suprathermal_toroidal_velocity_factor + * Exp6D( + -0.5 + * ( + ((r_6d - r_peak) ** 2 / toroidal_velocity_peak_width_R**2) + + ((z6d - z_peak) ** 2 / toroidal_velocity_peak_width_Z**2) + ) + ) +) + +# Convert supra-thermal toroidal velocity to Cartesian components +vx_suprathermal_6d = -v_toroidal_suprathermal_6d * y6d / (r_6d + 1e-10) +vy_suprathermal_6d = v_toroidal_suprathermal_6d * x6d / (r_6d + 1e-10) +vz_suprathermal_6d = 0.0 + +factor_st = (deuterium_mass / (2 * pi * elementary_charge * suprathermal_temperature)) ** 1.5 + +suprathermal_exponential_6d = Exp6D( + -0.5 + * deuterium_mass + * ( + (vx6d - vx_suprathermal_6d) ** 2 + + (vy6d - vy_suprathermal_6d) ** 2 + + (vz6d - vz_suprathermal_6d) ** 2 + ) + / (elementary_charge * suprathermal_temperature) +) +supra_thermal_pdf_6d = ( + supra_thermal_population_ratio * density_6d * factor_st * suprathermal_exponential_6d +) +phase_space_density_6d = bulk_pdf_6d + supra_thermal_pdf_6d + + +generic_distribution = GenericDistribution( + phase_space_density_6d, density_3d, temperature_3d, bulk_velocity_profile +) + +# Example evaluation: sample the distribution at a point in R-Z space +# At R=2, Z=0 (the peak), sample velocity distribution +# Convert R=2, Z=0 to Cartesian: x=2, y=0, z=0 +sample_x, sample_y, sample_z = 2.0, 0.0, 0.0 + +# Sample velocity distribution along the toroidal and vertical directions (vy and vz components) +v_vals = np.linspace(-1e6, 3e6, 1000) +n_particles_vy = np.zeros_like(v_vals) +n_particles_vz = np.zeros_like(v_vals) +for i in range(len(v_vals)): + n_particles_vy[i] = generic_distribution(sample_x, sample_y, sample_z, 0.0, v_vals[i], 0.0) + n_particles_vz[i] = generic_distribution(sample_x, sample_y, sample_z, 0.0, 0.0, v_vals[i]) + + +_, ax = plt.subplots() +ax.plot(v_vals, n_particles_vy, label="$\\mathrm{v}_\\mathrm{y}$") +ax.plot(v_vals, n_particles_vz, label="$\\mathrm{v}_\\mathrm{z}$") +ax.legend() +ax.set_xlabel("Velocity (m/s)") +ax.set_ylabel("Phase Space Density (s^3/m^6)") +ax.set_title("Velocity Distribution at R=2, Z=0 (Peak Location)") +ax.grid(True) + +# sample the temperature distribution in the R-Z plane +r_vals = np.linspace(1, 3, 100) +z_vals = np.linspace(-2, 2, 210) + +temperature_vals = np.zeros((r_vals.size, z_vals.size)) +density_vals = np.zeros((r_vals.size, z_vals.size)) +bulk_velocity_vals = np.zeros((r_vals.size, z_vals.size)) +for i in range(r_vals.size): + for j in range(z_vals.size): + temperature_vals[i, j] = generic_distribution.effective_temperature(r_vals[i], 0.0, z_vals[j]) + density_vals[i, j] = generic_distribution.density(r_vals[i], 0.0, z_vals[j]) + vector_velocity = generic_distribution.bulk_velocity(r_vals[i], 0.0, z_vals[j]) + bulk_velocity_vals[i, j] = np.sqrt(vector_velocity.x**2 + vector_velocity.y**2 + vector_velocity.z**2) + +_, ax = plt.subplots() +pcm = ax.pcolormesh(r_vals, z_vals, temperature_vals.transpose(), shading='gouraud') +plt.colorbar(pcm, ax=ax, label="Temperature (eV)") +ax.set_xlabel("R (m)") +ax.set_ylabel("Z (m)") +ax.set_title("Temperature Distribution") +ax.set_aspect('equal') +ax.grid(True) + +_, ax = plt.subplots() +pcm = ax.pcolormesh(r_vals, z_vals, density_vals.transpose(), shading='gouraud') +plt.colorbar(pcm, ax=ax, label="Density (m^-3)") +ax.set_xlabel("R (m)") +ax.set_ylabel("Z (m)") +ax.set_title("Density Distribution") +ax.set_aspect('equal') +ax.grid(True) + +_, ax = plt.subplots() +pcm = ax.pcolormesh(r_vals, z_vals, bulk_velocity_vals.transpose(), shading='gouraud') +plt.colorbar(pcm, ax=ax, label="velocity (m/s)") +ax.set_xlabel("R (m)") +ax.set_ylabel("Z (m)") +ax.set_title("Toroidal Bulk Velocity Distribution") +ax.set_aspect('equal') +ax.grid(True) + +# sample the x, y velocity components in the X-Y plane using arrow vectors +x_vals = np.linspace(-3, 3, 21) # Reduced resolution for clearer quiver plot +y_vals = np.linspace(-3, 3, 21) + +X, Y = np.meshgrid(x_vals, y_vals) +x_velocity_vals = np.zeros_like(X) +y_velocity_vals = np.zeros_like(Y) + +for i in range(x_vals.size): + for j in range(y_vals.size): + vector_velocity = generic_distribution.bulk_velocity(x_vals[i], y_vals[j], 0.0) + x_velocity_vals[j, i] = vector_velocity.x + y_velocity_vals[j, i] = vector_velocity.y + +# Calculate velocity magnitude for colormap +velocity_magnitude = np.sqrt(x_velocity_vals**2 + y_velocity_vals**2) + +_, ax = plt.subplots() +quiver = ax.quiver(X, Y, x_velocity_vals, y_velocity_vals, velocity_magnitude, + cmap='viridis', scale=1e6, width=0.003) +plt.colorbar(quiver, ax=ax, label="Velocity Magnitude (m/s)") +ax.set_xlabel("X (m)") +ax.set_ylabel("Y (m)") +ax.set_title("Bulk x, y Velocity Cmponents Field in X-Y Plane (Z=0)") +ax.set_aspect('equal') +ax.grid(True)