PySDM.backends.thrust_rtc

GPU-resident backend using NVRTC runtime compilation library for CUDA

 1"""
 2GPU-resident backend using NVRTC runtime compilation library for CUDA
 3"""
 4
 5import os
 6import warnings
 7
 8import numpy as np
 9
10from PySDM.backends.impl_thrust_rtc.conf import trtc
11from PySDM.backends.impl_thrust_rtc.methods.collisions_methods import CollisionsMethods
12from PySDM.backends.impl_thrust_rtc.methods.condensation_methods import (
13    CondensationMethods,
14)
15from PySDM.backends.impl_thrust_rtc.methods.displacement_methods import (
16    DisplacementMethods,
17)
18from PySDM.backends.impl_thrust_rtc.methods.freezing_methods import FreezingMethods
19from PySDM.backends.impl_thrust_rtc.methods.index_methods import IndexMethods
20from PySDM.backends.impl_thrust_rtc.methods.isotope_methods import IsotopeMethods
21from PySDM.backends.impl_thrust_rtc.methods.moments_methods import MomentsMethods
22from PySDM.backends.impl_thrust_rtc.methods.pair_methods import PairMethods
23from PySDM.backends.impl_thrust_rtc.methods.physics_methods import PhysicsMethods
24from PySDM.backends.impl_thrust_rtc.methods.terminal_velocity_methods import (
25    TerminalVelocityMethods,
26)
27from PySDM.backends.impl_thrust_rtc.random import Random as ImportedRandom
28from PySDM.backends.impl_thrust_rtc.storage import make_storage_class
29from PySDM.formulae import Formulae
30
31
32class ThrustRTC(  # pylint: disable=duplicate-code,too-many-ancestors
33    CollisionsMethods,
34    PairMethods,
35    IndexMethods,
36    PhysicsMethods,
37    CondensationMethods,
38    MomentsMethods,
39    DisplacementMethods,
40    TerminalVelocityMethods,
41    FreezingMethods,
42    IsotopeMethods,
43):
44    ENABLE = True
45    Random = ImportedRandom
46
47    default_croupier = "global"
48
49    def __init__(
50        self, formulae=None, double_precision=False, debug=False, verbose=False
51    ):
52        self.formulae = formulae or Formulae()
53
54        self._conv_function = trtc.DVDouble if double_precision else trtc.DVFloat
55        self._real_type = "double" if double_precision else "float"
56        self._np_dtype = np.float64 if double_precision else np.float32
57
58        self.Storage = make_storage_class(self)
59
60        CollisionsMethods.__init__(self)
61        PairMethods.__init__(self)
62        IndexMethods.__init__(self)
63        PhysicsMethods.__init__(self)
64        CondensationMethods.__init__(self)
65        MomentsMethods.__init__(self, double_precision=double_precision)
66        DisplacementMethods.__init__(self)
67        TerminalVelocityMethods.__init__(self)
68        FreezingMethods.__init__(self)
69
70        trtc.Set_Kernel_Debug(debug)
71        trtc.Set_Verbose(verbose)
72
73        if not ThrustRTC.ENABLE and "CI" not in os.environ:
74            warnings.warn("CUDA is not available, using FakeThrustRTC!")
33class ThrustRTC(  # pylint: disable=duplicate-code,too-many-ancestors
34    CollisionsMethods,
35    PairMethods,
36    IndexMethods,
37    PhysicsMethods,
38    CondensationMethods,
39    MomentsMethods,
40    DisplacementMethods,
41    TerminalVelocityMethods,
42    FreezingMethods,
43    IsotopeMethods,
44):
45    ENABLE = True
46    Random = ImportedRandom
47
48    default_croupier = "global"
49
50    def __init__(
51        self, formulae=None, double_precision=False, debug=False, verbose=False
52    ):
53        self.formulae = formulae or Formulae()
54
55        self._conv_function = trtc.DVDouble if double_precision else trtc.DVFloat
56        self._real_type = "double" if double_precision else "float"
57        self._np_dtype = np.float64 if double_precision else np.float32
58
59        self.Storage = make_storage_class(self)
60
61        CollisionsMethods.__init__(self)
62        PairMethods.__init__(self)
63        IndexMethods.__init__(self)
64        PhysicsMethods.__init__(self)
65        CondensationMethods.__init__(self)
66        MomentsMethods.__init__(self, double_precision=double_precision)
67        DisplacementMethods.__init__(self)
68        TerminalVelocityMethods.__init__(self)
69        FreezingMethods.__init__(self)
70
71        trtc.Set_Kernel_Debug(debug)
72        trtc.Set_Verbose(verbose)
73
74        if not ThrustRTC.ENABLE and "CI" not in os.environ:
75            warnings.warn("CUDA is not available, using FakeThrustRTC!")
ThrustRTC(formulae=None, double_precision=False, debug=False, verbose=False)
50    def __init__(
51        self, formulae=None, double_precision=False, debug=False, verbose=False
52    ):
53        self.formulae = formulae or Formulae()
54
55        self._conv_function = trtc.DVDouble if double_precision else trtc.DVFloat
56        self._real_type = "double" if double_precision else "float"
57        self._np_dtype = np.float64 if double_precision else np.float32
58
59        self.Storage = make_storage_class(self)
60
61        CollisionsMethods.__init__(self)
62        PairMethods.__init__(self)
63        IndexMethods.__init__(self)
64        PhysicsMethods.__init__(self)
65        CondensationMethods.__init__(self)
66        MomentsMethods.__init__(self, double_precision=double_precision)
67        DisplacementMethods.__init__(self)
68        TerminalVelocityMethods.__init__(self)
69        FreezingMethods.__init__(self)
70
71        trtc.Set_Kernel_Debug(debug)
72        trtc.Set_Verbose(verbose)
73
74        if not ThrustRTC.ENABLE and "CI" not in os.environ:
75            warnings.warn("CUDA is not available, using FakeThrustRTC!")
ENABLE = False
default_croupier = 'global'
formulae
Storage
Inherited Members
PySDM.backends.impl_thrust_rtc.methods.collisions_methods.CollisionsMethods
adaptive_sdm_end
scale_prob_for_adaptive_sdm_gamma
cell_id
collision_coalescence
collision_coalescence_breakup
compute_gamma
make_cell_caretaker
normalize
remove_zero_n_or_flagged
exp_fragmentation
gauss_fragmentation
slams_fragmentation
feingold1988_fragmentation
straub_fragmentation
PySDM.backends.impl_thrust_rtc.methods.pair_methods.PairMethods
distance_pair
find_pairs
max_pair
sort_pair
sort_within_pair_by_attr
sum_pair
min_pair
multiply_pair
PySDM.backends.impl_thrust_rtc.methods.index_methods.IndexMethods
identity_index
shuffle_global
shuffle_local
sort_by_key
PySDM.backends.impl_thrust_rtc.methods.physics_methods.PhysicsMethods
critical_volume
temperature_pressure_rh
explicit_euler
volume_of_water_mass
mass_of_water_volume
air_density
air_dynamic_viscosity
reynolds_number
PySDM.backends.impl_thrust_rtc.methods.condensation_methods.CondensationMethods
keys
RH_rtol
adaptive
max_iters
ml_old
ml_new
T
dthd_dt_pred
d_water_vapour_mixing_ratio__dt_predicted
drhod_dt_pred
rhod_copy
m_d
vars
vars_data
calculate_m_l
condensation
make_condensation_solver
PySDM.backends.impl_thrust_rtc.methods.moments_methods.MomentsMethods
ensure_floating_point
moments
spectrum_moments
PySDM.backends.impl_thrust_rtc.methods.displacement_methods.DisplacementMethods
calculate_displacement
flag_precipitated
flag_out_of_column
PySDM.backends.impl_thrust_rtc.methods.terminal_velocity_methods.TerminalVelocityMethods
linear_collection_efficiency
interpolation
power_series
terminal_velocity
PySDM.backends.impl_thrust_rtc.methods.freezing_methods.FreezingMethods
freeze_time_dependent_body
freeze_singular_body
freeze_singular
freeze_time_dependent
PySDM.backends.impl_thrust_rtc.methods.isotope_methods.IsotopeMethods
isotopic_delta
isotopic_fractionation
class ThrustRTC.Random(PySDM.backends.impl_common.random_common.RandomCommon):
65    class Random(RandomCommon):  # pylint: disable=too-few-public-methods
66        def __init__(self, size, seed):
67            super().__init__(size, seed)
68            self.generator = np.random.default_rng(seed)
69
70        def __call__(self, storage):
71            # pylint: disable=unsupported-assignment-operation
72            storage.data.ndarray[:] = self.generator.uniform(0, 1, storage.shape)