PySDM.backends.thrust_rtc

GPU-resident backend using NVRTC runtime compilation library for CUDA

View Source

 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) View Source

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: immersion_freezing_time_dependent_body; immersion_freezing_singular_body; thaw_instantaneous_body; homogeneous_freezing_threshold_body; homogeneous_freezing_time_dependent_body; immersion_freezing_singular; immersion_freezing_time_dependent; thaw_instantaneous; homogeneous_freezing_threshold; homogeneous_freezing_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): View Source

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)

Inherited Members

PySDM.backends.Random: Random; generator
PySDM.backends.impl_common.random_common.RandomCommon: size