Skip to content

AbsorptionCuda

Bases: CudaContribution

Computes the contribution to the optical depth occuring from molecular absorption.

Source code in src\taurex_cupy\contributions\absorption.py
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
class AbsorptionCuda(CudaContribution):
    """
    Computes the contribution to the optical depth
    occuring from molecular absorption.
    """

    def __init__(self):
        super().__init__("Absorption")
        self._opacity_cache: CudaOpacityCache = CudaOpacityCache()
        self._xsec_cache = {}

    def build(self, model: OneDForwardModel):
        super().build(model)
        self._opacity_cache.set_native_grid(model.nativeWavenumberGrid)

    def prepare_each(
        self, model: OneDForwardModel, wngrid: npt.NDArray[np.floating]
    ) -> t.Iterator[tuple[str, npt.NDArray[np.floating]]]:
        """
        Prepares each molecular opacity by weighting them
        by their mixing ratio in the atmosphere

        Parameters
        ----------
        model: :class:`~taurex.model.model.ForwardModel`
            Forward model

        wngrid: :obj:`array`
            Wavenumber grid

        Yields
        ------
        component: :obj:`tuple` of type (str, :obj:`array`)
            Name of molecule and weighted opacity

        """

        self.debug("Preparing model with %s", wngrid.shape)
        self._ngrid = wngrid.shape[0]
        # Loop through all active gases
        for gas in model.chemistry.activeGases:
            # Get the mix ratio of the gas
            gas_mix = model.chemistry.get_gas_mix_profile(gas)
            self.info("Recomputing active gas %s opacity", gas)

            # Get the cross section object relating to the gas
            xsec: CudaOpacity = self._opacity_cache[gas]
            sigma_xsec = xsec.opacity(
                model.temperatureProfile,
                model.pressureProfile,
                gas_mix,
                wngrid=wngrid,
            )

            # Temporarily assign to master cross-section
            self.sigma_xsec = sigma_xsec
            yield gas, sigma_xsec
        sigma_xsec = None

    @classmethod
    def input_keywords(cls):
        return [
            "AbsorptionCuda",
        ]

prepare_each(model, wngrid)

Prepares each molecular opacity by weighting them by their mixing ratio in the atmosphere

Parameters

model: :class:~taurex.model.model.ForwardModel Forward model

:obj:array

Wavenumber grid

Yields

component: :obj:tuple of type (str, :obj:array) Name of molecule and weighted opacity

Source code in src\taurex_cupy\contributions\absorption.py
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
def prepare_each(
    self, model: OneDForwardModel, wngrid: npt.NDArray[np.floating]
) -> t.Iterator[tuple[str, npt.NDArray[np.floating]]]:
    """
    Prepares each molecular opacity by weighting them
    by their mixing ratio in the atmosphere

    Parameters
    ----------
    model: :class:`~taurex.model.model.ForwardModel`
        Forward model

    wngrid: :obj:`array`
        Wavenumber grid

    Yields
    ------
    component: :obj:`tuple` of type (str, :obj:`array`)
        Name of molecule and weighted opacity

    """

    self.debug("Preparing model with %s", wngrid.shape)
    self._ngrid = wngrid.shape[0]
    # Loop through all active gases
    for gas in model.chemistry.activeGases:
        # Get the mix ratio of the gas
        gas_mix = model.chemistry.get_gas_mix_profile(gas)
        self.info("Recomputing active gas %s opacity", gas)

        # Get the cross section object relating to the gas
        xsec: CudaOpacity = self._opacity_cache[gas]
        sigma_xsec = xsec.opacity(
            model.temperatureProfile,
            model.pressureProfile,
            gas_mix,
            wngrid=wngrid,
        )

        # Temporarily assign to master cross-section
        self.sigma_xsec = sigma_xsec
        yield gas, sigma_xsec
    sigma_xsec = None

CIACuda

Bases: CudaContribution

Source code in src\taurex_cupy\contributions\cia.py
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
class CIACuda(CudaContribution):
    def __init__(self, cia_pairs=None):
        super().__init__("CIA")
        self._opacity_cache = CudaCiaCache()
        self._xsec_cache = {}
        self._cia_pairs = cia_pairs
        if self._cia_pairs is None:
            self._cia_pairs = []

    def build(self, model):
        super().build(model)
        self._opacity_cache.set_native_grid(model.nativeWavenumberGrid)

    def contribute(
        self,
        model,
        start_layer,
        end_layer,
        density_offset,
        layer,
        density,
        tau,
        path_length=None,
        with_layer_offset: bool = False,
    ):
        """
        Computes an integral for a single layer for the optical depth.

        Parameters
        ----------
        model: :class:`~taurex.model.model.ForwardModel`
            A forward model

        start_layer: int
            Lowest layer limit for integration

        end_layer: int
            Upper layer limit of integration

        density_offset: int
            offset in density layer

        layer: int
            atmospheric layer being computed

        density: :obj:`array`
            density profile of atmosphere

        tau: :obj:`array`
            optical depth to store result

        path_length: :obj:`array`
            integration length

        """

        self.debug(
            " %s %s %s %s %s %s %s",
            start_layer,
            end_layer,
            density_offset,
            layer,
            density,
            tau,
            self._ngrid,
        )

        cuda_contribute_tau(
            start_layer,
            end_layer,
            density_offset,
            self.sigma_xsec,
            density * density,
            path_length,
            self._nlayers,
            self._ngrid,
            tau,
            with_layer_offset=with_layer_offset,
        )
        self.debug("DONE")

    @property
    def ciaPairs(self):
        """
        Returns list of molecular pairs involved

        Returns
        -------
        :obj:`list` of str
        """

        return self._cia_pairs

    @ciaPairs.setter
    def ciaPairs(self, value):
        self._cia_pairs = value

    def prepare_each(
        self, model: OneDForwardModel, wngrid: npt.NDArray[np.floating]
    ) -> t.Iterator[tuple[str, npt.NDArray[np.floating]]]:
        """
        Prepares each molecular opacity by weighting them
        by their mixing ratio in the atmosphere

        Parameters
        ----------
        model: :class:`~taurex.model.model.ForwardModel`
            Forward model

        wngrid: :obj:`array`
            Wavenumber grid

        Yields
        ------
        component: :obj:`tuple` of type (str, :obj:`array`)
            Name of molecule and weighted opacity

        """

        self.debug("Preparing model with %s", wngrid.shape)
        self._ngrid = wngrid.shape[0]
        self._nlayers = model.nLayers
        sigma_xsec = cp.zeros(
            shape=(model.nLayers, wngrid.shape[0]),
            dtype=np.float64,
        )

        chemistry = model.chemistry
        # Loop through all active gases
        for pairName in self.ciaPairs:
            xsec = self._opacity_cache[pairName]
            cia = xsec._xsec
            cia_factor = cp.array(
                chemistry.get_gas_mix_profile(cia.pairOne) * chemistry.get_gas_mix_profile(cia.pairTwo)
            )

            # Get the cross section object relating to the gas

            sigma_xsec = xsec.opacity(model.temperatureProfile, cia_factor, wngrid=wngrid)

            # Temporarily assign to master cross-section
            self.sigma_xsec = sigma_xsec
            yield pairName, sigma_xsec

    def write(self, output):
        contrib = super().write(output)
        if len(self.ciaPairs) > 0:
            contrib.write_string_array("cia_pairs", self.ciaPairs)
        return contrib

    @classmethod
    def input_keywords(cls):
        return [
            "CIACuda",
        ]

ciaPairs property writable

Returns list of molecular pairs involved

Returns

:obj:list of str

contribute(model, start_layer, end_layer, density_offset, layer, density, tau, path_length=None, with_layer_offset=False)

Computes an integral for a single layer for the optical depth.

Parameters

model: :class:~taurex.model.model.ForwardModel A forward model

int

Lowest layer limit for integration

int

Upper layer limit of integration

int

offset in density layer

int

atmospheric layer being computed

:obj:array

density profile of atmosphere

:obj:array

optical depth to store result

:obj:array

integration length

Source code in src\taurex_cupy\contributions\cia.py
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
def contribute(
    self,
    model,
    start_layer,
    end_layer,
    density_offset,
    layer,
    density,
    tau,
    path_length=None,
    with_layer_offset: bool = False,
):
    """
    Computes an integral for a single layer for the optical depth.

    Parameters
    ----------
    model: :class:`~taurex.model.model.ForwardModel`
        A forward model

    start_layer: int
        Lowest layer limit for integration

    end_layer: int
        Upper layer limit of integration

    density_offset: int
        offset in density layer

    layer: int
        atmospheric layer being computed

    density: :obj:`array`
        density profile of atmosphere

    tau: :obj:`array`
        optical depth to store result

    path_length: :obj:`array`
        integration length

    """

    self.debug(
        " %s %s %s %s %s %s %s",
        start_layer,
        end_layer,
        density_offset,
        layer,
        density,
        tau,
        self._ngrid,
    )

    cuda_contribute_tau(
        start_layer,
        end_layer,
        density_offset,
        self.sigma_xsec,
        density * density,
        path_length,
        self._nlayers,
        self._ngrid,
        tau,
        with_layer_offset=with_layer_offset,
    )
    self.debug("DONE")

prepare_each(model, wngrid)

Prepares each molecular opacity by weighting them by their mixing ratio in the atmosphere

Parameters

model: :class:~taurex.model.model.ForwardModel Forward model

:obj:array

Wavenumber grid

Yields

component: :obj:tuple of type (str, :obj:array) Name of molecule and weighted opacity

Source code in src\taurex_cupy\contributions\cia.py
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
def prepare_each(
    self, model: OneDForwardModel, wngrid: npt.NDArray[np.floating]
) -> t.Iterator[tuple[str, npt.NDArray[np.floating]]]:
    """
    Prepares each molecular opacity by weighting them
    by their mixing ratio in the atmosphere

    Parameters
    ----------
    model: :class:`~taurex.model.model.ForwardModel`
        Forward model

    wngrid: :obj:`array`
        Wavenumber grid

    Yields
    ------
    component: :obj:`tuple` of type (str, :obj:`array`)
        Name of molecule and weighted opacity

    """

    self.debug("Preparing model with %s", wngrid.shape)
    self._ngrid = wngrid.shape[0]
    self._nlayers = model.nLayers
    sigma_xsec = cp.zeros(
        shape=(model.nLayers, wngrid.shape[0]),
        dtype=np.float64,
    )

    chemistry = model.chemistry
    # Loop through all active gases
    for pairName in self.ciaPairs:
        xsec = self._opacity_cache[pairName]
        cia = xsec._xsec
        cia_factor = cp.array(
            chemistry.get_gas_mix_profile(cia.pairOne) * chemistry.get_gas_mix_profile(cia.pairTwo)
        )

        # Get the cross section object relating to the gas

        sigma_xsec = xsec.opacity(model.temperatureProfile, cia_factor, wngrid=wngrid)

        # Temporarily assign to master cross-section
        self.sigma_xsec = sigma_xsec
        yield pairName, sigma_xsec

CudaContribution

Bases: Contribution

Source code in src\taurex_cupy\contributions\cudacontrib.py
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
class CudaContribution(Contribution):
    def __init__(self, name):
        super().__init__(name)
        self._is_cuda_model = False

    def contribute(
        self,
        model: OneDForwardModel,
        start_layer: cp.ndarray,
        end_layer: cp.ndarray,
        density_offset: cp.ndarray,
        layer: int,
        density: cp.ndarray,
        tau: cp.ndarray,
        path_length: t.Optional[cp.ndarray] = None,
        with_layer_offset: bool = True,
    ):
        """
        Computes an integral for a single layer for the optical depth.

        Parameters
        ----------
        model: :class:`~taurex.model.model.ForwardModel`
            A forward model

        start_layer: int
            Lowest layer limit for integration

        end_layer: int
            Upper layer limit of integration

        density_offset: int
            offset in density layer

        layer: int
            atmospheric layer being computed

        density: :obj:`array`
            density profile of atmosphere

        tau: :obj:`array`
            optical depth to store result

        path_length: :obj:`array`
            integration length

        """
        self.debug(
            " %s %s %s %s %s %s %s",
            start_layer,
            end_layer,
            density_offset,
            layer,
            density,
            tau,
            self._ngrid,
        )

        cuda_contribute_tau(
            start_layer,
            end_layer,
            density_offset,
            self.sigma_xsec,
            density,
            path_length,
            self._nlayers,
            self._ngrid,
            tau,
            with_layer_offset=with_layer_offset,
        )

        self.debug("DONE")

    def prepare(self, model, wngrid):
        """

        Used to prepare the contribution for the calculation.
        Called before the forward model performs the main optical depth
        calculation. Default behaviour is to loop through :func:`prepare_each`
        and sum all results into a single cross-section.

        Parameters
        ----------
        model: :class:`~taurex.model.model.ForwardModel`
            Forward model

        wngrid: :obj:`array`
            Wavenumber grid
        """
        del self.sigma_xsec
        self._ngrid = wngrid.shape[0]
        self._nlayers = model.nLayers

        sigma_xsec = cp.zeros(shape=(self._nlayers, self._ngrid), dtype=np.float64)

        for gas, sigma in self.prepare_each(model, wngrid):
            self.debug("Gas %s", gas)
            self.debug("Sigma %s", sigma)
            sigma_xsec += sigma

        self.sigma_xsec = sigma_xsec
        self.debug("Final sigma is %s", self.sigma_xsec)
        self.info("Done")

contribute(model, start_layer, end_layer, density_offset, layer, density, tau, path_length=None, with_layer_offset=True)

Computes an integral for a single layer for the optical depth.

Parameters

model: :class:~taurex.model.model.ForwardModel A forward model

int

Lowest layer limit for integration

int

Upper layer limit of integration

int

offset in density layer

int

atmospheric layer being computed

:obj:array

density profile of atmosphere

:obj:array

optical depth to store result

:obj:array

integration length

Source code in src\taurex_cupy\contributions\cudacontrib.py
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
def contribute(
    self,
    model: OneDForwardModel,
    start_layer: cp.ndarray,
    end_layer: cp.ndarray,
    density_offset: cp.ndarray,
    layer: int,
    density: cp.ndarray,
    tau: cp.ndarray,
    path_length: t.Optional[cp.ndarray] = None,
    with_layer_offset: bool = True,
):
    """
    Computes an integral for a single layer for the optical depth.

    Parameters
    ----------
    model: :class:`~taurex.model.model.ForwardModel`
        A forward model

    start_layer: int
        Lowest layer limit for integration

    end_layer: int
        Upper layer limit of integration

    density_offset: int
        offset in density layer

    layer: int
        atmospheric layer being computed

    density: :obj:`array`
        density profile of atmosphere

    tau: :obj:`array`
        optical depth to store result

    path_length: :obj:`array`
        integration length

    """
    self.debug(
        " %s %s %s %s %s %s %s",
        start_layer,
        end_layer,
        density_offset,
        layer,
        density,
        tau,
        self._ngrid,
    )

    cuda_contribute_tau(
        start_layer,
        end_layer,
        density_offset,
        self.sigma_xsec,
        density,
        path_length,
        self._nlayers,
        self._ngrid,
        tau,
        with_layer_offset=with_layer_offset,
    )

    self.debug("DONE")

prepare(model, wngrid)

Used to prepare the contribution for the calculation. Called before the forward model performs the main optical depth calculation. Default behaviour is to loop through :func:prepare_each and sum all results into a single cross-section.

Parameters

model: :class:~taurex.model.model.ForwardModel Forward model

:obj:array

Wavenumber grid

Source code in src\taurex_cupy\contributions\cudacontrib.py
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
def prepare(self, model, wngrid):
    """

    Used to prepare the contribution for the calculation.
    Called before the forward model performs the main optical depth
    calculation. Default behaviour is to loop through :func:`prepare_each`
    and sum all results into a single cross-section.

    Parameters
    ----------
    model: :class:`~taurex.model.model.ForwardModel`
        Forward model

    wngrid: :obj:`array`
        Wavenumber grid
    """
    del self.sigma_xsec
    self._ngrid = wngrid.shape[0]
    self._nlayers = model.nLayers

    sigma_xsec = cp.zeros(shape=(self._nlayers, self._ngrid), dtype=np.float64)

    for gas, sigma in self.prepare_each(model, wngrid):
        self.debug("Gas %s", gas)
        self.debug("Sigma %s", sigma)
        sigma_xsec += sigma

    self.sigma_xsec = sigma_xsec
    self.debug("Final sigma is %s", self.sigma_xsec)
    self.info("Done")

DirectImageCudaModel

Bases: EmissionCudaModel

A forward model for direct imaging of exo-planets.

Source code in src\taurex_cupy\model\directimage.py
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
class DirectImageCudaModel(EmissionCudaModel):
    """A forward model for direct imaging of exo-planets."""

    def compute_final_flux(self, f_total: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
        """Compute the final flux.

        This is the emission flux that is observed at the telescope directly
        from an exo-planet.

        """
        return compute_direct_image_final_flux(f_total, self._planet.fullRadius, self._star.distance * 3.08567758e16)

    @classmethod
    def input_keywords(cls) -> tuple[str, ...]:
        """Input keywords for this class."""
        return (
            "direct_cuda",
            "directimage_cuda",
        )

compute_final_flux(f_total)

Compute the final flux.

This is the emission flux that is observed at the telescope directly from an exo-planet.

Source code in src\taurex_cupy\model\directimage.py
13
14
15
16
17
18
19
20
def compute_final_flux(self, f_total: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
    """Compute the final flux.

    This is the emission flux that is observed at the telescope directly
    from an exo-planet.

    """
    return compute_direct_image_final_flux(f_total, self._planet.fullRadius, self._star.distance * 3.08567758e16)

input_keywords() classmethod

Input keywords for this class.

Source code in src\taurex_cupy\model\directimage.py
22
23
24
25
26
27
28
@classmethod
def input_keywords(cls) -> tuple[str, ...]:
    """Input keywords for this class."""
    return (
        "direct_cuda",
        "directimage_cuda",
    )

EmissionCudaModel

Bases: SimpleForwardModel

A forward model for eclipse models using CUDA

Parameters

:class:~taurex.data.planet.Planet, optional

Planet model, default planet is Jupiter

:class:~taurex.data.stellar.star.Star, optional

Star model, default star is Sun-like

:class:~taurex.data.profiles.pressure.pressureprofile.PressureProfile, optional

Pressure model, alternative is to set nlayers, atm_min_pressure and atm_max_pressure

:class:~taurex.data.profiles.temperature.tprofile.TemperatureProfile, optional

Temperature model, default is an :class:~taurex.data.profiles.temperature.isothermal.Isothermal profile at 1500 K

:class:~taurex.data.profiles.chemistry.chemistry.Chemistry, optional

Chemistry model, default is :class:~taurex.data.profiles.chemistry.taurexchemistry.TaurexChemistry with H2O and CH4

int, optional

Number of layers. Used if pressure_profile is not defined.

float, optional

Pressure at TOA. Used if pressure_profile is not defined.

float, optional

Pressure at BOA. Used if pressure_profile is not defined.

int, optional

Number of Gaussian quadrature points. Default is 4

Source code in src\taurex_cupy\model\eclipse.py
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
class EmissionCudaModel(SimpleForwardModel):
    """

    A forward model for eclipse models using CUDA

    Parameters
    ----------

    planet: :class:`~taurex.data.planet.Planet`, optional
        Planet model, default planet is Jupiter

    star: :class:`~taurex.data.stellar.star.Star`, optional
        Star model, default star is Sun-like

    pressure_profile: :class:`~taurex.data.profiles.pressure.pressureprofile.PressureProfile`, optional
        Pressure model, alternative is to set ``nlayers``, ``atm_min_pressure``
        and ``atm_max_pressure``

    temperature_profile: :class:`~taurex.data.profiles.temperature.tprofile.TemperatureProfile`, optional
        Temperature model, default is an :class:`~taurex.data.profiles.temperature.isothermal.Isothermal`
        profile at 1500 K

    chemistry: :class:`~taurex.data.profiles.chemistry.chemistry.Chemistry`, optional
        Chemistry model, default is
        :class:`~taurex.data.profiles.chemistry.taurexchemistry.TaurexChemistry` with
        ``H2O`` and ``CH4``

    nlayers: int, optional
        Number of layers. Used if ``pressure_profile`` is not defined.

    atm_min_pressure: float, optional
        Pressure at TOA. Used if ``pressure_profile`` is not defined.

    atm_max_pressure: float, optional
        Pressure at BOA. Used if ``pressure_profile`` is not defined.

    ngauss: int, optional
        Number of Gaussian quadrature points. Default is 4

    """

    def __init__(
        self,
        planet=None,
        star=None,
        pressure_profile=None,
        temperature_profile=None,
        chemistry=None,
        nlayers=100,
        atm_min_pressure=1e-4,
        atm_max_pressure=1e6,
        ngauss=4,
    ):
        super().__init__(
            self.__class__.__name__,
            planet,
            star,
            pressure_profile,
            temperature_profile,
            chemistry,
            nlayers,
            atm_min_pressure,
            atm_max_pressure,
        )

        self.set_num_gauss(ngauss)

    def set_num_gauss(self, value):
        self._ngauss = int(value)
        mu, weight = np.polynomial.legendre.leggauss(self._ngauss)
        self._mu_quads = (mu + 1) / 2
        self._wi_quads = (weight) / 2

        self._gpu_mu_quads = cp.array(self._mu_quads)
        self._gpu_wi_quads = cp.array(self._wi_quads)

    def build(self):
        super().build()
        for contrib in self.contribution_list:
            contrib.build(self)
        self._start_surface_K = cp.array(np.array([0]).astype(np.int32))
        self._end_surface_K = cp.array(np.array([self.nLayers]).astype(np.int32))

        self._start_layer = cp.arange(self.nLayers, dtype=np.int32) + 1
        self._end_layer = cp.full_like(self._start_layer, self.nLayers)

        self._start_dtau = cp.arange(self.nLayers)
        self._end_dtau = self._start_dtau + 1

        self._dz = cp.zeros(
            shape=(
                self.nLayers,
                self.nLayers,
            ),
            dtype=np.float64,
        )
        self._density_offset = cp.zeros(shape=(self.nLayers,), dtype=np.int32)

        # self._tau_buffer= drv.pagelocked_zeros(shape=(self.nativeWavenumberGrid.shape[-1], self.nLayers,),dtype=np.float64)

    def partial_model(self, wngrid=None, cutoff_grid=True):
        from taurex.util.util import clip_native_to_wngrid

        self.initialize_profiles()

        native_grid = self.nativeWavenumberGrid
        if wngrid is not None and cutoff_grid:
            native_grid = clip_native_to_wngrid(native_grid, wngrid)
        self._star.initialize(native_grid)

        for contrib in self.contribution_list:
            contrib.prepare(self, native_grid)

        return self.evaluate_emission(native_grid, False)

    @property
    def cuda_contributions(self):
        return [c for c in self.contribution_list if isinstance(c, CudaContribution)]

    @property
    def non_cuda_contributions(self):
        return [c for c in self.contribution_list if not isinstance(c, CudaContribution)]

    def evaluate_emission(self, wngrid, return_contrib, keep_in_gpu=False):
        total_layers = self.nLayers

        dz = self.deltaz

        dz = np.array([dz for x in range(self.nLayers)])
        self._dz = cp.array(dz)

        wngrid_size = wngrid.shape[0]
        temperature = self.temperatureProfile
        density_profile = cp.array(self.densityProfile)

        self._fully_cuda = len(self.non_cuda_contributions) == 0

        layer_tau = cp.zeros(
            shape=(total_layers, wngrid_size),
            dtype=np.float64,
        )
        dtau = cp.zeros(
            shape=(total_layers, wngrid_size),
            dtype=np.float64,
        )

        intensity = cp.zeros(
            shape=(self._ngauss, wngrid_size),
            dtype=np.float64,
        )
        blackbody = cuda_blackbody(wngrid, temperature.ravel())
        tau_host = cpx.zeros_pinned(shape=(total_layers, wngrid_size), dtype=np.float64)
        if not self._fully_cuda:
            self.fallback_noncuda(layer_tau, dtau, wngrid, total_layers)
        for contrib in self.cuda_contributions:
            contrib.contribute(
                self,
                self._start_layer,
                self._end_layer,
                self._density_offset,
                0,
                density_profile,
                layer_tau,
                path_length=self._dz,
                with_layer_offset=False,
            )
            contrib.contribute(
                self,
                self._start_dtau,
                self._end_dtau,
                self._density_offset,
                0,
                density_profile,
                dtau,
                path_length=self._dz,
                with_layer_offset=False,
            )

        integral_kernal = gen_partial_kernal(self._ngauss, self.nLayers, wngrid_size)

        THREAD_PER_BLOCK_X = 64

        NUM_BLOCK_X = int(math.ceil(wngrid_size / THREAD_PER_BLOCK_X))

        integral_kernal(
            (NUM_BLOCK_X,),
            (THREAD_PER_BLOCK_X,),
            (intensity, layer_tau, dtau, blackbody),
        )
        layer_tau.get(out=tau_host)
        if keep_in_gpu:
            return (
                intensity,
                1 / self._gpu_mu_quads[:, None],
                self._gpu_wi_quads[:, None],
                tau_host,
            )

        return (
            intensity.get(),
            1 / self._mu_quads[:, None],
            self._wi_quads[:, None],
            tau_host,
        )

    def path_integral(self, wngrid, return_contrib):
        intensity, _mu, _w, tau = self.evaluate_emission(wngrid, return_contrib, keep_in_gpu=True)
        self.debug("I: %s", intensity)

        flux_total = 2.0 * np.pi * cp.sum(intensity * _w / _mu, axis=0)
        self.debug("flux_total %s", flux_total)

        return self.compute_final_flux(flux_total).ravel().get(), tau

    def fallback_noncuda(self, gpu_layer_tau, gpu_dtau, wngrid, total_layers):
        wngrid_size = wngrid.shape[0]

        dz = np.zeros(total_layers)
        dz[:-1] = np.diff(self.altitudeProfile)
        dz[-1] = self.altitudeProfile[-1] - self.altitudeProfile[-2]

        density = self.densityProfile
        layer_tau = np.zeros(shape=(total_layers, wngrid_size))
        dtau = np.zeros(shape=(total_layers, wngrid_size))
        _dtau = np.zeros(shape=(1, wngrid_size))
        _layer_tau = np.zeros(shape=(1, wngrid_size))

        # Loop upwards
        for layer in range(total_layers):
            _layer_tau[...] = 0.0
            _dtau[...] = 0.0
            for contrib in self.non_cuda_contributions:
                contrib.contribute(
                    self,
                    layer + 1,
                    total_layers,
                    0,
                    0,
                    density,
                    _layer_tau,
                    path_length=dz,
                )
                contrib.contribute(self, layer, layer + 1, 0, 0, density, _dtau, path_length=dz)

            layer_tau[layer, :] += _layer_tau[0]
            dtau[layer, :] += _dtau[0]

        gpu_layer_tau.set(layer_tau)
        gpu_dtau.set(dtau)

    def compute_final_flux(self, f_total):
        star_sed = cp.array(self._star.spectralEmissionDensity)

        self.debug("Star SED: %s", star_sed)
        # quit()
        star_radius = self._star.radius
        planet_radius = self._planet.fullRadius
        self.debug("star_radius %s", self._star.radius)
        self.debug("planet_radius %s", self._star.radius)
        last_flux = (f_total / star_sed) * (planet_radius / star_radius) ** 2

        self.debug("last_flux %s", last_flux)

        return last_flux

    @classmethod
    def input_keywords(cls) -> tuple[str, ...]:
        return (
            "emission_cuda",
            "eclipse_cuda",
        )

FlatMieCuda

Bases: CudaContribution

Computes a flat (gray) absorption contribution.

Absorption is computed as a flat value between two pressures across all wavenumbers.

Parameters

float

Opacity value

float

Bottom of absorbing region in Pa

float

Top of absorbing region in Pa

Source code in src\taurex_cupy\contributions\flatmie.py
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
class FlatMieCuda(CudaContribution):
    """Computes a flat (gray) absorption contribution.

    Absorption is computed as a flat value between two pressures
    across all wavenumbers.

    Parameters
    ----------

    flat_mix_ratio: float
        Opacity value

    flat_bottomP: float
        Bottom of absorbing region in Pa

    flat_topP: float
        Top of absorbing region in Pa

    """

    def __init__(
        self,
        flat_mix_ratio: float = 1e-10,
        flat_bottomP: float = -1,
        flat_topP: float = -1,
    ) -> None:
        super().__init__("Mie")

        self._mie_mix = flat_mix_ratio
        self._mie_bottom_pressure = flat_bottomP
        self._mie_top_pressure = flat_topP

    @fitparam(
        param_name="flat_topP",
        param_latex=r"$P^{mie}_\mathrm{top}$",
        default_mode="log",
        default_fit=False,
        default_bounds=[1e-20, 1],
    )
    def mieTopPressure(self) -> float:
        """
        Pressure at top of absorbing region in Pa
        """
        return self._mie_top_pressure

    @mieTopPressure.setter
    def mieTopPressure(self, value: float) -> None:
        self._mie_top_pressure = value

    @fitparam(
        param_name="flat_bottomP",
        param_latex=r"$P^{mie}_\mathrm{bottom}$",
        default_mode="log",
        default_fit=False,
        default_bounds=[1e-20, 1],
    )
    def mieBottomPressure(self) -> float:
        """Pressure at bottom of absorbing region in Pa."""
        return self._mie_bottom_pressure

    @mieBottomPressure.setter
    def mieBottomPressure(self, value: float) -> None:
        self._mie_bottom_pressure = value

    @fitparam(
        param_name="flat_mix_ratio",
        param_latex=r"$\chi_\mathrm{mie}$",
        default_mode="log",
        default_fit=False,
        default_bounds=[1e-20, 1],
    )
    def mieMixing(self) -> float:
        """Opacity of absorbing region in :math:`m^2`."""
        return self._mie_mix

    @mieMixing.setter
    def mieMixing(self, value: float) -> None:
        self._mie_mix = value

    def prepare_each(
        self, model: OneDForwardModel, wngrid: npt.NDArray[np.float64]
    ) -> t.Generator[tuple[str, npt.NDArray[np.float64]], None, None]:
        """Computes and flat absorbing opacity for the pressure regions given.


        Parameters
        ----------
        model: :class:`~taurex.model.model.ForwardModel`
            Forward model

        wngrid: :obj:`array`
            Wavenumber grid

        Yields
        ------
        component: :obj:`tuple` of type (str, :obj:`array`)
            ``Flat`` and the weighted mie opacity.


        """
        self._nlayers = model.nLayers
        self._ngrid = wngrid.shape[0]

        pressure_levels = np.log10(model.pressure.pressure_profile_levels[::-1])

        bottom_pressure = self.mieBottomPressure
        if bottom_pressure < 0:
            bottom_pressure = pressure_levels.max()

        top_pressure = np.log10(self.mieTopPressure)
        if top_pressure < 0:
            top_pressure = pressure_levels.min()

        p_left = pressure_levels[:-1]
        p_right = pressure_levels[1:]

        p_range = sorted([top_pressure, bottom_pressure])

        save_start = np.searchsorted(p_right, p_range[0], side="right")
        save_stop = np.searchsorted(p_left[1:], p_range[1], side="right")
        p_min = p_left[save_start : save_stop + 1]
        p_max = p_right[save_start : save_stop + 1]
        weight = np.minimum(p_range[-1], p_max) - np.maximum(p_range[0], p_min)
        weight /= weight.max()
        gpu_weight = cp.asarray(weight, dtype=np.float64)
        sigma_xsec = cp.zeros(shape=(self._nlayers, wngrid.shape[0]))
        sigma_xsec[save_start : save_stop + 1] = gpu_weight[:, None] * self.mieMixing

        sigma_xsec = sigma_xsec[::-1]

        self.sigma_xsec = sigma_xsec

        yield "Flat", sigma_xsec

    def write(self, output: OutputGroup) -> OutputGroup:
        """Write contribution to output.

        Parameters
        ----------
        output: :class:`~taurex.output.output.Output`
            Output object to write to

        Returns
        -------
        output: :class:`~taurex.output.output.Output`
            Output object that was written to

        """
        contrib = super().write(output)
        contrib.write_scalar("flat_mix_ratio", self._mie_mix)
        contrib.write_scalar("flat_bottomP", self._mie_bottom_pressure)
        contrib.write_scalar("flat_topP", self._mie_top_pressure)
        return contrib

    @classmethod
    def input_keywords(cls) -> tuple[str]:
        """Return input keywords for the contribution."""
        return ("FlatMieCuda",)

input_keywords() classmethod

Return input keywords for the contribution.

Source code in src\taurex_cupy\contributions\flatmie.py
169
170
171
172
@classmethod
def input_keywords(cls) -> tuple[str]:
    """Return input keywords for the contribution."""
    return ("FlatMieCuda",)

prepare_each(model, wngrid)

Computes and flat absorbing opacity for the pressure regions given.

Parameters

model: :class:~taurex.model.model.ForwardModel Forward model

:obj:array

Wavenumber grid

Yields

component: :obj:tuple of type (str, :obj:array) Flat and the weighted mie opacity.

Source code in src\taurex_cupy\contributions\flatmie.py
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
def prepare_each(
    self, model: OneDForwardModel, wngrid: npt.NDArray[np.float64]
) -> t.Generator[tuple[str, npt.NDArray[np.float64]], None, None]:
    """Computes and flat absorbing opacity for the pressure regions given.


    Parameters
    ----------
    model: :class:`~taurex.model.model.ForwardModel`
        Forward model

    wngrid: :obj:`array`
        Wavenumber grid

    Yields
    ------
    component: :obj:`tuple` of type (str, :obj:`array`)
        ``Flat`` and the weighted mie opacity.


    """
    self._nlayers = model.nLayers
    self._ngrid = wngrid.shape[0]

    pressure_levels = np.log10(model.pressure.pressure_profile_levels[::-1])

    bottom_pressure = self.mieBottomPressure
    if bottom_pressure < 0:
        bottom_pressure = pressure_levels.max()

    top_pressure = np.log10(self.mieTopPressure)
    if top_pressure < 0:
        top_pressure = pressure_levels.min()

    p_left = pressure_levels[:-1]
    p_right = pressure_levels[1:]

    p_range = sorted([top_pressure, bottom_pressure])

    save_start = np.searchsorted(p_right, p_range[0], side="right")
    save_stop = np.searchsorted(p_left[1:], p_range[1], side="right")
    p_min = p_left[save_start : save_stop + 1]
    p_max = p_right[save_start : save_stop + 1]
    weight = np.minimum(p_range[-1], p_max) - np.maximum(p_range[0], p_min)
    weight /= weight.max()
    gpu_weight = cp.asarray(weight, dtype=np.float64)
    sigma_xsec = cp.zeros(shape=(self._nlayers, wngrid.shape[0]))
    sigma_xsec[save_start : save_stop + 1] = gpu_weight[:, None] * self.mieMixing

    sigma_xsec = sigma_xsec[::-1]

    self.sigma_xsec = sigma_xsec

    yield "Flat", sigma_xsec

write(output)

Write contribution to output.

Parameters

output: :class:~taurex.output.output.Output Output object to write to

Returns

output: :class:~taurex.output.output.Output Output object that was written to

Source code in src\taurex_cupy\contributions\flatmie.py
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
def write(self, output: OutputGroup) -> OutputGroup:
    """Write contribution to output.

    Parameters
    ----------
    output: :class:`~taurex.output.output.Output`
        Output object to write to

    Returns
    -------
    output: :class:`~taurex.output.output.Output`
        Output object that was written to

    """
    contrib = super().write(output)
    contrib.write_scalar("flat_mix_ratio", self._mie_mix)
    contrib.write_scalar("flat_bottomP", self._mie_bottom_pressure)
    contrib.write_scalar("flat_topP", self._mie_top_pressure)
    return contrib

LeeMieCuda

Bases: CudaContribution

Computes Mie scattering contribution to optica depth.

Formalism taken from: Lee et al. 2013, ApJ, 778, 97

Parameters

float

Particle radius in um

float

Extinction coefficient

float

Mixing ratio in atmosphere

float

Bottom of cloud deck in Pa

float

Top of cloud deck in Pa

Source code in src\taurex_cupy\contributions\leemie.py
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
class LeeMieCuda(CudaContribution):
    """Computes Mie scattering contribution to optica depth.

    Formalism taken from: Lee et al. 2013, ApJ, 778, 97

    Parameters
    ----------

    lee_mie_radius: float
        Particle radius in um

    lee_mie_q: float
        Extinction coefficient

    lee_mie_mix_ratio: float
        Mixing ratio in atmosphere

    lee_mie_bottomP: float
        Bottom of cloud deck in Pa

    lee_mie_topP: float
        Top of cloud deck in Pa


    """

    def __init__(
        self,
        lee_mie_radius: t.Optional[float] = 0.01,
        lee_mie_q: t.Optional[float] = 40,
        lee_mie_mix_ratio: t.Optional[float] = 1e-10,
        lee_mie_bottomP: t.Optional[float] = -1,
        lee_mie_topP: t.Optional[float] = -1,
    ) -> None:
        super().__init__("Mie")

        self._mie_radius = lee_mie_radius
        self._mie_q = lee_mie_q
        self._mie_mix = lee_mie_mix_ratio
        self._mie_bottom_pressure = lee_mie_bottomP
        self._mie_top_pressure = lee_mie_topP

    @fitparam(
        param_name="lee_mie_radius",
        param_latex=r"$R^{lee}_{\mathrm{mie}}$",
        default_fit=False,
        default_bounds=[0.01, 0.5],
    )
    def mieRadius(self) -> float:
        """Particle radius in um."""
        return self._mie_radius

    @mieRadius.setter
    def mieRadius(self, value: float) -> None:
        """Particle radius in um."""
        self._mie_radius = value

    @fitparam(
        param_name="lee_mie_q",
        param_latex=r"$Q_\mathrm{ext}$",
        default_fit=False,
        default_bounds=[-10, 1],
    )
    def mieQ(self) -> float:
        """Extinction coefficient."""
        return self._mie_q

    @mieQ.setter
    def mieQ(self, value: float) -> None:
        self._mie_q = value

    @fitparam(
        param_name="lee_mie_topP",
        param_latex=r"$P^{lee}_\mathrm{top}$",
        default_mode="log",
        default_fit=False,
        default_bounds=[-1, 1],
    )
    def mieTopPressure(self) -> float:
        """Pressure at top of cloud deck in Pa."""
        return self._mie_top_pressure

    @mieTopPressure.setter
    def mieTopPressure(self, value: float) -> None:
        """Pressure at top of cloud deck in Pa."""
        self._mie_top_pressure = value

    @fitparam(
        param_name="lee_mie_bottomP",
        param_latex=r"$P^{lee}_\mathrm{bottom}$",
        default_mode="log",
        default_fit=False,
        default_bounds=[-1, 1],
    )
    def mieBottomPressure(self) -> float:
        """Pressure at bottom of cloud deck in Pa."""
        return self._mie_bottom_pressure

    @mieBottomPressure.setter
    def mieBottomPressure(self, value: float) -> None:
        """Pressure at bottom of cloud deck in Pa."""
        self._mie_bottom_pressure = value

    @fitparam(
        param_name="lee_mie_mix_ratio",
        param_latex=r"$\chi^{lee}_\mathrm{mie}$",
        default_mode="log",
        default_fit=False,
        default_bounds=[-1, 1],
    )
    def mieMixing(self) -> float:
        """Mixing ratio in atmosphere."""
        return self._mie_mix

    @mieMixing.setter
    def mieMixing(self, value: float) -> None:
        """Mixing ratio in atmosphere."""
        self._mie_mix = value

    def prepare_each(
        self, model: OneDForwardModel, wngrid: npt.NDArray[np.float64]
    ) -> t.Generator[tuple[str, npt.NDArray[np.float64]], None, None]:
        """Compute and weights the mie opacity for the pressure regions given.

        Parameters
        ----------
        model: :class:`~taurex.model.model.ForwardModel`
            Forward model

        wngrid: :obj:`array`
            Wavenumber grid

        Yields
        ------
        component: :obj:`tuple` of type (str, :obj:`array`)
            ``Lee`` and the weighted mie opacity.

        """
        self._nlayers = model.nLayers
        self._ngrid = wngrid.shape[0]

        pressure_profile = cp.asarray(model.pressureProfile)

        bottom_pressure = self.mieBottomPressure
        if bottom_pressure < 0:
            bottom_pressure = model.pressureProfile[0]

        top_pressure = self.mieTopPressure
        if top_pressure < 0:
            top_pressure = model.pressureProfile[-1]

        wltmp = cp.asarray(10000 / wngrid)

        a = self.mieRadius

        x = 2.0 * np.pi * a / wltmp
        self.debug("wngrid %s", wngrid)
        self.debug("x %s", x)
        q_ext = 5.0 / (self.mieQ * x ** (-4.0) + x ** (0.2))

        sigma_xsec = cp.zeros(shape=(self._nlayers, wngrid.shape[0]))

        # This must transform um to the xsec format in TauREx (m2)
        am = a * 1e-6

        sigma_mie = q_ext * np.pi * (am**2.0)

        self.debug("q_ext %s", q_ext)
        self.debug("radius um %s", a)
        self.debug("sigma %s", sigma_mie)

        self.debug("bottome_pressure %s", bottom_pressure)
        self.debug("top_pressure %s", top_pressure)

        cloud_filter = (pressure_profile <= bottom_pressure) & (pressure_profile >= top_pressure)

        sigma_xsec[cloud_filter, ...] = sigma_mie * self.mieMixing

        self.sigma_xsec = sigma_xsec

        self.debug("final xsec %s", self.sigma_xsec)

        yield "Lee", sigma_xsec

    def write(self, output: OutputGroup) -> OutputGroup:
        """Write output group."""
        contrib = super().write(output)
        contrib.write_scalar("lee_mie_radius", self._mie_radius)
        contrib.write_scalar("lee_mie_q", self._mie_q)
        contrib.write_scalar("lee_mie_mix_ratio", self._mie_mix)
        contrib.write_scalar("lee_mie_bottomP", self._mie_bottom_pressure)
        contrib.write_scalar("lee_mie_topP", self._mie_top_pressure)
        return contrib

    @classmethod
    def input_keywords(cls) -> tuple[str]:
        """Input keywords.""" ""
        return ("LeeMieCuda",)

    BIBTEX_ENTRIES = LeeMieContribution.BIBTEX_ENTRIES

input_keywords() classmethod

Input keywords.

Source code in src\taurex_cupy\contributions\leemie.py
210
211
212
213
@classmethod
def input_keywords(cls) -> tuple[str]:
    """Input keywords.""" ""
    return ("LeeMieCuda",)

mieBottomPressure(value)

Pressure at bottom of cloud deck in Pa.

Source code in src\taurex_cupy\contributions\leemie.py
114
115
116
117
@mieBottomPressure.setter
def mieBottomPressure(self, value: float) -> None:
    """Pressure at bottom of cloud deck in Pa."""
    self._mie_bottom_pressure = value

mieMixing(value)

Mixing ratio in atmosphere.

Source code in src\taurex_cupy\contributions\leemie.py
130
131
132
133
@mieMixing.setter
def mieMixing(self, value: float) -> None:
    """Mixing ratio in atmosphere."""
    self._mie_mix = value

mieRadius(value)

Particle radius in um.

Source code in src\taurex_cupy\contributions\leemie.py
68
69
70
71
@mieRadius.setter
def mieRadius(self, value: float) -> None:
    """Particle radius in um."""
    self._mie_radius = value

mieTopPressure(value)

Pressure at top of cloud deck in Pa.

Source code in src\taurex_cupy\contributions\leemie.py
 98
 99
100
101
@mieTopPressure.setter
def mieTopPressure(self, value: float) -> None:
    """Pressure at top of cloud deck in Pa."""
    self._mie_top_pressure = value

prepare_each(model, wngrid)

Compute and weights the mie opacity for the pressure regions given.

Parameters

model: :class:~taurex.model.model.ForwardModel Forward model

:obj:array

Wavenumber grid

Yields

component: :obj:tuple of type (str, :obj:array) Lee and the weighted mie opacity.

Source code in src\taurex_cupy\contributions\leemie.py
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
def prepare_each(
    self, model: OneDForwardModel, wngrid: npt.NDArray[np.float64]
) -> t.Generator[tuple[str, npt.NDArray[np.float64]], None, None]:
    """Compute and weights the mie opacity for the pressure regions given.

    Parameters
    ----------
    model: :class:`~taurex.model.model.ForwardModel`
        Forward model

    wngrid: :obj:`array`
        Wavenumber grid

    Yields
    ------
    component: :obj:`tuple` of type (str, :obj:`array`)
        ``Lee`` and the weighted mie opacity.

    """
    self._nlayers = model.nLayers
    self._ngrid = wngrid.shape[0]

    pressure_profile = cp.asarray(model.pressureProfile)

    bottom_pressure = self.mieBottomPressure
    if bottom_pressure < 0:
        bottom_pressure = model.pressureProfile[0]

    top_pressure = self.mieTopPressure
    if top_pressure < 0:
        top_pressure = model.pressureProfile[-1]

    wltmp = cp.asarray(10000 / wngrid)

    a = self.mieRadius

    x = 2.0 * np.pi * a / wltmp
    self.debug("wngrid %s", wngrid)
    self.debug("x %s", x)
    q_ext = 5.0 / (self.mieQ * x ** (-4.0) + x ** (0.2))

    sigma_xsec = cp.zeros(shape=(self._nlayers, wngrid.shape[0]))

    # This must transform um to the xsec format in TauREx (m2)
    am = a * 1e-6

    sigma_mie = q_ext * np.pi * (am**2.0)

    self.debug("q_ext %s", q_ext)
    self.debug("radius um %s", a)
    self.debug("sigma %s", sigma_mie)

    self.debug("bottome_pressure %s", bottom_pressure)
    self.debug("top_pressure %s", top_pressure)

    cloud_filter = (pressure_profile <= bottom_pressure) & (pressure_profile >= top_pressure)

    sigma_xsec[cloud_filter, ...] = sigma_mie * self.mieMixing

    self.sigma_xsec = sigma_xsec

    self.debug("final xsec %s", self.sigma_xsec)

    yield "Lee", sigma_xsec

write(output)

Write output group.

Source code in src\taurex_cupy\contributions\leemie.py
200
201
202
203
204
205
206
207
208
def write(self, output: OutputGroup) -> OutputGroup:
    """Write output group."""
    contrib = super().write(output)
    contrib.write_scalar("lee_mie_radius", self._mie_radius)
    contrib.write_scalar("lee_mie_q", self._mie_q)
    contrib.write_scalar("lee_mie_mix_ratio", self._mie_mix)
    contrib.write_scalar("lee_mie_bottomP", self._mie_bottom_pressure)
    contrib.write_scalar("lee_mie_topP", self._mie_top_pressure)
    return contrib

RayleighCuda

Bases: CudaContribution

Computes the contribution to the optical depth occuring from molecular Rayleigh.

Source code in src\taurex_cupy\contributions\rayleigh.py
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
class RayleighCuda(CudaContribution):
    """
    Computes the contribution to the optical depth
    occuring from molecular Rayleigh.
    """

    def __init__(self):
        super().__init__("Rayleigh")
        self.sigmas = {}
        self._current_grid = None

    def build(self, model: OneDForwardModel):
        super().build(model)
        self._mix_array = cp.zeros(shape=(model.nLayers,), dtype=np.float64)
        self._current_grid = None
        self.sigmas.clear()

    def prepare_each(
        self, model: OneDForwardModel, wngrid: npt.NDArray[np.floating]
    ) -> t.Iterator[tuple[str, npt.NDArray[np.floating]]]:
        """
        Prepares each molecular opacity by weighting them
        by their mixing ratio in the atmosphere

        Parameters
        ----------
        model: :class:`~taurex.model.model.ForwardModel`
            Forward model

        wngrid: :obj:`array`
            Wavenumber grid

        Yields
        ------
        component: :obj:`tuple` of type (str, :obj:`array`)
            Name of molecule and weighted opacity

        """
        from taurex.util.scattering import rayleigh_sigma_from_name

        self.debug("Preparing model with %s", wngrid.shape)
        self._ngrid = wngrid.shape[0]
        molecules = model.chemistry.activeGases + model.chemistry.inactiveGases
        if self._current_grid is None or not np.array_equal(self._current_grid, wngrid):
            gpu_wngrid = cp.asarray(wngrid, dtype=np.float64)
            for gasname in molecules:
                self.sigmas[gasname] = rayleigh_sigma_from_name(gasname, gpu_wngrid)

        for gasname in molecules:
            if np.max(model.chemistry.get_gas_mix_profile(gasname)) == 0.0:
                continue
            sigma = self.sigmas.get(gasname, None)

            if sigma is not None:
                final_sigma = sigma[None, :] * cp.array(model.chemistry.get_gas_mix_profile(gasname)[:, None])
                self.sigma_xsec = final_sigma
                yield gasname, final_sigma

    @classmethod
    def input_keywords(cls):
        return [
            "RayleighCuda",
        ]

    BIBTEX_ENTRIES = RayleighContribution.BIBTEX_ENTRIES

prepare_each(model, wngrid)

Prepares each molecular opacity by weighting them by their mixing ratio in the atmosphere

Parameters

model: :class:~taurex.model.model.ForwardModel Forward model

:obj:array

Wavenumber grid

Yields

component: :obj:tuple of type (str, :obj:array) Name of molecule and weighted opacity

Source code in src\taurex_cupy\contributions\rayleigh.py
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
def prepare_each(
    self, model: OneDForwardModel, wngrid: npt.NDArray[np.floating]
) -> t.Iterator[tuple[str, npt.NDArray[np.floating]]]:
    """
    Prepares each molecular opacity by weighting them
    by their mixing ratio in the atmosphere

    Parameters
    ----------
    model: :class:`~taurex.model.model.ForwardModel`
        Forward model

    wngrid: :obj:`array`
        Wavenumber grid

    Yields
    ------
    component: :obj:`tuple` of type (str, :obj:`array`)
        Name of molecule and weighted opacity

    """
    from taurex.util.scattering import rayleigh_sigma_from_name

    self.debug("Preparing model with %s", wngrid.shape)
    self._ngrid = wngrid.shape[0]
    molecules = model.chemistry.activeGases + model.chemistry.inactiveGases
    if self._current_grid is None or not np.array_equal(self._current_grid, wngrid):
        gpu_wngrid = cp.asarray(wngrid, dtype=np.float64)
        for gasname in molecules:
            self.sigmas[gasname] = rayleigh_sigma_from_name(gasname, gpu_wngrid)

    for gasname in molecules:
        if np.max(model.chemistry.get_gas_mix_profile(gasname)) == 0.0:
            continue
        sigma = self.sigmas.get(gasname, None)

        if sigma is not None:
            final_sigma = sigma[None, :] * cp.array(model.chemistry.get_gas_mix_profile(gasname)[:, None])
            self.sigma_xsec = final_sigma
            yield gasname, final_sigma

SimpleCloudsCuda

Bases: CudaContribution

Optically thick cloud deck up to a certain height

These have the form:

.. math:: \tau(\lambda,z) = \begin{cases} \infty & \quad \text{if } P(z) >= P_{0}\ 0 & \quad \text{if } P(z) < P_{0} \end{cases}

Where :math:P_{0} is the pressure at the top of the cloud-deck.

Source code in src\taurex_cupy\contributions\simpleclouds.py
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
class SimpleCloudsCuda(CudaContribution):
    """
    Optically thick cloud deck up to a certain height

    These have the form:

    .. math::
            \\tau(\\lambda,z) =
                \\begin{cases}
                \\infty       & \\quad \\text{if } P(z) >= P_{0}\\\\
                0            & \\quad \\text{if } P(z) < P_{0}
                \\end{cases}

    Where :math:`P_{0}` is the pressure at the top of the cloud-deck.


    """

    def __init__(self, clouds_pressure: t.Optional[float] = 1e3) -> None:
        """Initialize the cloud model.

        Parameters
        ----------
        clouds_pressure : float, optional
            Pressure at top of cloud deck, by default 1e3


        """
        super().__init__("SimpleClouds")
        self._cloud_pressure = clouds_pressure

    @property
    def order(self) -> int:
        return 3

    def contribute(
        self,
        model: OneDForwardModel,
        start_layer: int,
        end_layer: int,
        density_offset: int,
        layer: int,
        density: npt.NDArray[np.float64],
        tau: npt.NDArray[np.float64],
        path_length: t.Optional[npt.NDArray[np.float64]] = None,
    ):
        """Contribute the cloud opacity to the optical depth."""
        tau[layer] += self.sigma_xsec[layer, :]

    def prepare_each(
        self, model: OneDForwardModel, wngrid: npt.NDArray[np.float64]
    ) -> t.Generator[tuple[str, npt.NDArray[np.float64]], None, None]:
        """Compute cross-section that is infinitely absorbing.


        Absorption occurs up to a certain height.

        Parameters
        ----------
        model: :class:`~taurex.model.model.ForwardModel`
            Forward model

        wngrid: :obj:`array`
            Wavenumber grid

        Yields
        ------
        component: :obj:`tuple` of type (str, :obj:`array`)
            ``Clouds`` and opacity array.


        """

        contrib = cp.zeros(
            shape=(
                model.nLayers,
                wngrid.shape[0],
            )
        )
        cloud_filtr = model.pressureProfile >= self._cloud_pressure
        contrib[cloud_filtr, :] = np.inf
        self._contrib = contrib
        yield "Clouds", self._contrib

    @fitparam(
        param_name="clouds_pressure",
        param_latex=r"$P_\mathrm{clouds}$",
        default_mode="log",
        default_fit=False,
        default_bounds=[1e-3, 1e6],
    )
    def cloudsPressure(self):
        """Cloud top pressure in Pascal."""
        return self._cloud_pressure

    @cloudsPressure.setter
    def cloudsPressure(self, value):
        """Cloud top pressure in Pascal."""
        self._cloud_pressure = value

    def write(self, output: OutputGroup):
        """Write the cloud pressure to the output."""
        contrib = super().write(output)
        contrib.write_scalar("clouds_pressure", self._cloud_pressure)
        return contrib

    @classmethod
    def input_keywords(cls) -> tuple[str, str]:
        return (
            "SimpleCloudsCuda",
            "ThickCloudsCuda",
        )

__init__(clouds_pressure=1000.0)

Initialize the cloud model.

Parameters

clouds_pressure : float, optional Pressure at top of cloud deck, by default 1e3

Source code in src\taurex_cupy\contributions\simpleclouds.py
33
34
35
36
37
38
39
40
41
42
43
44
def __init__(self, clouds_pressure: t.Optional[float] = 1e3) -> None:
    """Initialize the cloud model.

    Parameters
    ----------
    clouds_pressure : float, optional
        Pressure at top of cloud deck, by default 1e3


    """
    super().__init__("SimpleClouds")
    self._cloud_pressure = clouds_pressure

cloudsPressure(value)

Cloud top pressure in Pascal.

Source code in src\taurex_cupy\contributions\simpleclouds.py
110
111
112
113
@cloudsPressure.setter
def cloudsPressure(self, value):
    """Cloud top pressure in Pascal."""
    self._cloud_pressure = value

contribute(model, start_layer, end_layer, density_offset, layer, density, tau, path_length=None)

Contribute the cloud opacity to the optical depth.

Source code in src\taurex_cupy\contributions\simpleclouds.py
50
51
52
53
54
55
56
57
58
59
60
61
62
def contribute(
    self,
    model: OneDForwardModel,
    start_layer: int,
    end_layer: int,
    density_offset: int,
    layer: int,
    density: npt.NDArray[np.float64],
    tau: npt.NDArray[np.float64],
    path_length: t.Optional[npt.NDArray[np.float64]] = None,
):
    """Contribute the cloud opacity to the optical depth."""
    tau[layer] += self.sigma_xsec[layer, :]

prepare_each(model, wngrid)

Compute cross-section that is infinitely absorbing.

Absorption occurs up to a certain height.

Parameters

model: :class:~taurex.model.model.ForwardModel Forward model

:obj:array

Wavenumber grid

Yields

component: :obj:tuple of type (str, :obj:array) Clouds and opacity array.

Source code in src\taurex_cupy\contributions\simpleclouds.py
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
def prepare_each(
    self, model: OneDForwardModel, wngrid: npt.NDArray[np.float64]
) -> t.Generator[tuple[str, npt.NDArray[np.float64]], None, None]:
    """Compute cross-section that is infinitely absorbing.


    Absorption occurs up to a certain height.

    Parameters
    ----------
    model: :class:`~taurex.model.model.ForwardModel`
        Forward model

    wngrid: :obj:`array`
        Wavenumber grid

    Yields
    ------
    component: :obj:`tuple` of type (str, :obj:`array`)
        ``Clouds`` and opacity array.


    """

    contrib = cp.zeros(
        shape=(
            model.nLayers,
            wngrid.shape[0],
        )
    )
    cloud_filtr = model.pressureProfile >= self._cloud_pressure
    contrib[cloud_filtr, :] = np.inf
    self._contrib = contrib
    yield "Clouds", self._contrib

write(output)

Write the cloud pressure to the output.

Source code in src\taurex_cupy\contributions\simpleclouds.py
115
116
117
118
119
def write(self, output: OutputGroup):
    """Write the cloud pressure to the output."""
    contrib = super().write(output)
    contrib.write_scalar("clouds_pressure", self._cloud_pressure)
    return contrib

TransmissionCudaModel

Bases: OneDForwardModel

A forward model for transits using GPU acceleration.

Source code in src\taurex_cupy\model\transit.py
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
class TransmissionCudaModel(OneDForwardModel):
    """A forward model for transits using GPU acceleration."""

    def __init__(
        self,
        planet: t.Optional[Planet] = None,
        star: t.Optional[Star] = None,
        pressure_profile: t.Optional[PressureProfile] = None,
        temperature_profile: t.Optional[TemperatureProfile] = None,
        chemistry: t.Optional[Chemistry] = None,
        nlayers: t.Optional[int] = 100,
        atm_min_pressure: t.Optional[float] = 1e-4,
        atm_max_pressure: t.Optional[float] = 1e6,
        contributions: t.Optional[list[Contribution | CudaContribution]] = None,
    ):
        """Initialise the model.

        Args:
            planet: Planet object
            star: Star object
            pressure_profile: Pressure profile for the atmosphere
            temperature_profile: Temperature profile for the atmosphere
            chemistry: Chemical model
            nlayers: Number of layers
            atm_min_pressure: Minimum pressure (If pressure profile is not set)
            atm_max_pressure: Maximum pressure (If pressure profile is not set)
            contributions: List of contributions

        """
        super().__init__(
            name=self.__class__.__name__,
            planet=planet,
            star=star,
            pressure_profile=pressure_profile,
            temperature_profile=temperature_profile,
            chemistry=chemistry,
            nlayers=nlayers,
            atm_min_pressure=atm_min_pressure,
            atm_max_pressure=atm_max_pressure,
            contributions=contributions,
        )

    def compute_path_length(self, dz: npt.NDArray[np.floating]) -> list[npt.NDArray[np.float64]]:
        r"""Compute path length for each layer, new method.

        Args:
            dz: $\Delta z$ of the layer (altitude)

        Returns:
            list: Path length for each layer

        """
        from taurex.util.geometry import parallel_vector

        altitude_boundaries = self.altitude_boundaries
        radius = self.planet.fullRadius

        # Generate our line of sight paths
        viewer, tangent = parallel_vector(radius, self.altitude_profile + dz / 2, altitude_boundaries.max())

        path_lengths = self.planet.compute_path_length(altitude_boundaries, viewer, tangent)
        # We need to pad the path lengths to the number of layers
        # path_lengths = [l for _, l in path_lengths]
        dls = []
        for _, p in path_lengths:
            # We need to pad the path lengths to the number of layers
            p = np.pad(p, (0, self.nLayers - len(p)), "constant", constant_values=0)
            dls.append(p)

        return dls
        #

    @property
    def cuda_contributions(self) -> None:
        """Get contributions that use cuda."""
        return [c for c in self.contribution_list if isinstance(c, CudaContribution)]

    @property
    def non_cuda_contributions(self) -> None:
        """Get contributions that do not use cuda."""
        return [c for c in self.contribution_list if not isinstance(c, CudaContribution)]

    def build(self) -> None:
        """Build the model."""
        super().build()
        for contrib in self.contribution_list:
            contrib.build(self)
        self._startK = cp.array(np.array([0 for x in range(self.nLayers)]).astype(np.int32))
        self._endK = cp.array(np.array([self.nLayers - x for x in range(self.nLayers)]).astype(np.int32))
        self._density_offset = cp.array(np.array(list(range(self.nLayers))).astype(np.int32))

        # self._tau_buffer= drv.pagelocked_zeros(shape=(self.nativeWavenumberGrid.shape[-1], self.nLayers,),dtype=np.float64)

    def path_integral(
        self, wngrid: npt.NDArray[np.floating], return_contrib: bool
    ) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]:
        r"""Compute the path integral for the model.

        Args:
            wngrid: Wavenumber grid
            return_contrib: Return contributions

        Returns:
            tuple: $(R_p/R_s)^2$ and optical depth

        """
        total_layers = self.nLayers

        dz = self.deltaz

        wngrid_size = wngrid.shape[0]
        self._ngrid = wngrid_size
        cpu_dl = self.compute_path_length(dz)
        gpu_dl = cp.array(cpu_dl)
        density_profile = cp.array(self.densityProfile)

        self._fully_cuda = len(self.non_cuda_contributions) == 0

        tau = cp.zeros(
            shape=(total_layers, wngrid_size),
            dtype=np.float64,
        )

        tau_host = cpx.zeros_pinned(shape=(total_layers, wngrid_size), dtype=np.float64)
        if not self._fully_cuda:
            tau.set(self.fallback_noncuda(total_layers, cpu_dl, self.densityProfile, dz))

        for contrib in self.cuda_contributions:
            contrib.contribute(
                self,
                self._startK,
                self._endK,
                self._density_offset,
                0,
                density_profile,
                tau,
                path_length=gpu_dl,
            )

        rprs, tau = self.compute_absorption(tau, cp.array(dz))
        tau.get(out=tau_host)
        # cp.cuda.runtime.deviceSynchronize()
        final_rprs = rprs.get()

        return final_rprs, tau_host

    def fallback_noncuda(
        self,
        total_layers: int,
        path_length: npt.NDArray[np.floating],
        density_profile: npt.NDArray[np.floating],
        dz: npt.NDArray[np.floating],
    ) -> npt.NDArray[np.floating]:
        """Fallback for non-cuda contributions.

        This will compute them on the CPU before copying them to the GPU.

        Args:
            total_layers: Total layers
            path_length: Path length
            density_profile: Density profile
            dz: Delta altitude of the layer

        Returns:
            Optical depth

        """
        tau = np.zeros(shape=(total_layers, self._ngrid))
        for layer in range(total_layers):
            self.debug("Computing layer %s", layer)
            dl = path_length[layer]

            endK = total_layers - layer

            for contrib in self.non_cuda_contributions:
                self.debug("Adding contribution from %s", contrib.name)
                contrib.contribute(self, 0, endK, layer, layer, density_profile, tau, path_length=dl)
        return tau

    def compute_absorption(self, tau: cp.ndarray, dz: cp.ndarray) -> tuple[cp.ndarray, cp.ndarray]:
        r"""Compute the absorption.

        Args:
            tau: Tau
            dz: $\Delta z$ of the layer (altitude)

        Returns:
            tuple: $(R_p/R_s)^2$ and tau

        """
        cp.exp(-tau, out=tau)
        ap = cp.array(self.altitudeProfile[:, None])
        pradius = self._planet.fullRadius
        sradius = self._star.radius
        _dz = dz[:, None]

        integral = cp.sum((pradius + ap) * (1.0 - tau) * _dz * 2.0, axis=0)
        return ((pradius * pradius) + integral) / (sradius**2), tau

    @classmethod
    def input_keywords(cls):
        return [
            "transmission_cuda",
            "transit_cuda",
        ]

cuda_contributions property

Get contributions that use cuda.

non_cuda_contributions property

Get contributions that do not use cuda.

__init__(planet=None, star=None, pressure_profile=None, temperature_profile=None, chemistry=None, nlayers=100, atm_min_pressure=0.0001, atm_max_pressure=1000000.0, contributions=None)

Initialise the model.

Parameters:

Name Type Description Default
planet Optional[Planet]

Planet object

None
star Optional[Star]

Star object

None
pressure_profile Optional[PressureProfile]

Pressure profile for the atmosphere

None
temperature_profile Optional[TemperatureProfile]

Temperature profile for the atmosphere

None
chemistry Optional[Chemistry]

Chemical model

None
nlayers Optional[int]

Number of layers

100
atm_min_pressure Optional[float]

Minimum pressure (If pressure profile is not set)

0.0001
atm_max_pressure Optional[float]

Maximum pressure (If pressure profile is not set)

1000000.0
contributions Optional[list[Contribution | CudaContribution]]

List of contributions

None
Source code in src\taurex_cupy\model\transit.py
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
def __init__(
    self,
    planet: t.Optional[Planet] = None,
    star: t.Optional[Star] = None,
    pressure_profile: t.Optional[PressureProfile] = None,
    temperature_profile: t.Optional[TemperatureProfile] = None,
    chemistry: t.Optional[Chemistry] = None,
    nlayers: t.Optional[int] = 100,
    atm_min_pressure: t.Optional[float] = 1e-4,
    atm_max_pressure: t.Optional[float] = 1e6,
    contributions: t.Optional[list[Contribution | CudaContribution]] = None,
):
    """Initialise the model.

    Args:
        planet: Planet object
        star: Star object
        pressure_profile: Pressure profile for the atmosphere
        temperature_profile: Temperature profile for the atmosphere
        chemistry: Chemical model
        nlayers: Number of layers
        atm_min_pressure: Minimum pressure (If pressure profile is not set)
        atm_max_pressure: Maximum pressure (If pressure profile is not set)
        contributions: List of contributions

    """
    super().__init__(
        name=self.__class__.__name__,
        planet=planet,
        star=star,
        pressure_profile=pressure_profile,
        temperature_profile=temperature_profile,
        chemistry=chemistry,
        nlayers=nlayers,
        atm_min_pressure=atm_min_pressure,
        atm_max_pressure=atm_max_pressure,
        contributions=contributions,
    )

build()

Build the model.

Source code in src\taurex_cupy\model\transit.py
100
101
102
103
104
105
106
107
def build(self) -> None:
    """Build the model."""
    super().build()
    for contrib in self.contribution_list:
        contrib.build(self)
    self._startK = cp.array(np.array([0 for x in range(self.nLayers)]).astype(np.int32))
    self._endK = cp.array(np.array([self.nLayers - x for x in range(self.nLayers)]).astype(np.int32))
    self._density_offset = cp.array(np.array(list(range(self.nLayers))).astype(np.int32))

compute_absorption(tau, dz)

Compute the absorption.

Parameters:

Name Type Description Default
tau ndarray

Tau

required
dz ndarray

\(\Delta z\) of the layer (altitude)

required

Returns:

Name Type Description
tuple tuple[ndarray, ndarray]

\((R_p/R_s)^2\) and tau

Source code in src\taurex_cupy\model\transit.py
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
def compute_absorption(self, tau: cp.ndarray, dz: cp.ndarray) -> tuple[cp.ndarray, cp.ndarray]:
    r"""Compute the absorption.

    Args:
        tau: Tau
        dz: $\Delta z$ of the layer (altitude)

    Returns:
        tuple: $(R_p/R_s)^2$ and tau

    """
    cp.exp(-tau, out=tau)
    ap = cp.array(self.altitudeProfile[:, None])
    pradius = self._planet.fullRadius
    sradius = self._star.radius
    _dz = dz[:, None]

    integral = cp.sum((pradius + ap) * (1.0 - tau) * _dz * 2.0, axis=0)
    return ((pradius * pradius) + integral) / (sradius**2), tau

compute_path_length(dz)

Compute path length for each layer, new method.

Parameters:

Name Type Description Default
dz NDArray[floating]

\(\Delta z\) of the layer (altitude)

required

Returns:

Name Type Description
list list[NDArray[float64]]

Path length for each layer

Source code in src\taurex_cupy\model\transit.py
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
def compute_path_length(self, dz: npt.NDArray[np.floating]) -> list[npt.NDArray[np.float64]]:
    r"""Compute path length for each layer, new method.

    Args:
        dz: $\Delta z$ of the layer (altitude)

    Returns:
        list: Path length for each layer

    """
    from taurex.util.geometry import parallel_vector

    altitude_boundaries = self.altitude_boundaries
    radius = self.planet.fullRadius

    # Generate our line of sight paths
    viewer, tangent = parallel_vector(radius, self.altitude_profile + dz / 2, altitude_boundaries.max())

    path_lengths = self.planet.compute_path_length(altitude_boundaries, viewer, tangent)
    # We need to pad the path lengths to the number of layers
    # path_lengths = [l for _, l in path_lengths]
    dls = []
    for _, p in path_lengths:
        # We need to pad the path lengths to the number of layers
        p = np.pad(p, (0, self.nLayers - len(p)), "constant", constant_values=0)
        dls.append(p)

    return dls

fallback_noncuda(total_layers, path_length, density_profile, dz)

Fallback for non-cuda contributions.

This will compute them on the CPU before copying them to the GPU.

Parameters:

Name Type Description Default
total_layers int

Total layers

required
path_length NDArray[floating]

Path length

required
density_profile NDArray[floating]

Density profile

required
dz NDArray[floating]

Delta altitude of the layer

required

Returns:

Type Description
NDArray[floating]

Optical depth

Source code in src\taurex_cupy\model\transit.py
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
def fallback_noncuda(
    self,
    total_layers: int,
    path_length: npt.NDArray[np.floating],
    density_profile: npt.NDArray[np.floating],
    dz: npt.NDArray[np.floating],
) -> npt.NDArray[np.floating]:
    """Fallback for non-cuda contributions.

    This will compute them on the CPU before copying them to the GPU.

    Args:
        total_layers: Total layers
        path_length: Path length
        density_profile: Density profile
        dz: Delta altitude of the layer

    Returns:
        Optical depth

    """
    tau = np.zeros(shape=(total_layers, self._ngrid))
    for layer in range(total_layers):
        self.debug("Computing layer %s", layer)
        dl = path_length[layer]

        endK = total_layers - layer

        for contrib in self.non_cuda_contributions:
            self.debug("Adding contribution from %s", contrib.name)
            contrib.contribute(self, 0, endK, layer, layer, density_profile, tau, path_length=dl)
    return tau

path_integral(wngrid, return_contrib)

Compute the path integral for the model.

Parameters:

Name Type Description Default
wngrid NDArray[floating]

Wavenumber grid

required
return_contrib bool

Return contributions

required

Returns:

Name Type Description
tuple tuple[NDArray[floating], NDArray[floating]]

\((R_p/R_s)^2\) and optical depth

Source code in src\taurex_cupy\model\transit.py
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
def path_integral(
    self, wngrid: npt.NDArray[np.floating], return_contrib: bool
) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]:
    r"""Compute the path integral for the model.

    Args:
        wngrid: Wavenumber grid
        return_contrib: Return contributions

    Returns:
        tuple: $(R_p/R_s)^2$ and optical depth

    """
    total_layers = self.nLayers

    dz = self.deltaz

    wngrid_size = wngrid.shape[0]
    self._ngrid = wngrid_size
    cpu_dl = self.compute_path_length(dz)
    gpu_dl = cp.array(cpu_dl)
    density_profile = cp.array(self.densityProfile)

    self._fully_cuda = len(self.non_cuda_contributions) == 0

    tau = cp.zeros(
        shape=(total_layers, wngrid_size),
        dtype=np.float64,
    )

    tau_host = cpx.zeros_pinned(shape=(total_layers, wngrid_size), dtype=np.float64)
    if not self._fully_cuda:
        tau.set(self.fallback_noncuda(total_layers, cpu_dl, self.densityProfile, dz))

    for contrib in self.cuda_contributions:
        contrib.contribute(
            self,
            self._startK,
            self._endK,
            self._density_offset,
            0,
            density_profile,
            tau,
            path_length=gpu_dl,
        )

    rprs, tau = self.compute_absorption(tau, cp.array(dz))
    tau.get(out=tau_host)
    # cp.cuda.runtime.deviceSynchronize()
    final_rprs = rprs.get()

    return final_rprs, tau_host

Spectral module

Spectral functions for CUDA

cuda_blackbody(lamb, temperature)

Compute blackbody spectrum using cuda.

This will compute the blackbody spectrum using the formula:

\[ B_{\lambda} = \frac{2 \pi h c^2}{\lambda^5} \frac{1}{e^{\frac{hc}{\lambda k T}} - 1} \]

Parameters:

Name Type Description Default
lamb NDArray[floating]

Wavelength grid

required
temperature NDArray[floating]

Temperature grid

required

Returns:

Type Description
NDArray[floating]

Blackbody spectrum in W/m\(^2\)/micron/sr

Source code in src\taurex_cupy\spectral.py
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
def cuda_blackbody(lamb: npt.NDArray[np.floating], temperature: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]:
    r"""Compute blackbody spectrum using cuda.

    This will compute the blackbody spectrum using the formula:

    $$
    B_{\lambda} = \frac{2 \pi h c^2}{\lambda^5} \frac{1}{e^{\frac{hc}{\lambda k T}} - 1}
    $$

    Args:
        lamb: Wavelength grid
        temperature: Temperature grid

    Returns:
        Blackbody spectrum in W/m$^2$/micron/sr


    """
    from taurex.constants import KBOLTZ as k
    from taurex.constants import PLANCK as h
    from taurex.constants import SPDLIGT as c

    temperature = cp.atleast_1d(temperature)
    lamb = cp.atleast_1d(lamb)

    temperature = temperature[:, None]
    lamb = lamb[None, :]

    lamb = 1e-2 / lamb
    planck = 2 * np.pi * h * c**2 / lamb**5
    exp = cp.expm1(h * c / (lamb * k * temperature))
    bb = planck / (exp)

    return bb * 1e-6

Utility functions

utility functions for taurex-cupy

FakeCIA

Bases: CIA

Fake opacity for testing purposes.

Source code in src\taurex_cupy\util.py
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
class FakeCIA(CIA):
    """Fake opacity for testing purposes."""

    def __init__(
        self,
        molecule_pair: tuple[str, str],
        num_t: int = 27,
        wn_res: int = 15000,
        wn_size: tuple[float, float] = (300, 30000),
    ) -> None:
        """Create the fake opacity.

        Args:
            molecule_pair: The pair of molecules.
            num_t: The number of temperature points.
            wn_res: The resolution of the wavenumber grid.
            wn_size: The size of the wavenumber grid.



        """
        super().__init__("FAKE", "-".join(molecule_pair))
        self.pair = molecule_pair
        self._wavenumber_grid = create_grid_res(wn_res, *wn_size)[:, 0]
        self._temperature_grid = np.linspace(100, 10000, num_t)
        self._xsec_grid = np.random.rand(self._temperature_grid.size, self._wavenumber_grid.size)

    def find_closest_temperature_index(self, temperature: float) -> tuple[int, int]:
        """Finds the nearest indices for a particular temperature

        Args:
            temperature: The temperature to search for.

        Returns:
            tuple[int, int]: The indices of the closest temperatures.

        """

        t_min, t_max = find_closest_pair(self.temperatureGrid, temperature)
        return t_min, t_max

    def interp_linear_grid(self, temperature: float, t_idx_min: int, t_idx_max: int) -> npt.NDArray[np.float64]:
        """Linear interpolate the CIA opacity.

        Args:
            temperature: The temperature to interpolate.
            t_idx_min: The minimum temperature index.
            t_idx_max: The maximum temperature index.

        Returns:
            The interpolated opacity.


        """

        if temperature > self._temperature_grid.max():
            return self._xsec_grid[-1]
        elif temperature < self._temperature_grid.min():
            return self._xsec_grid[0]

        temp_max = self._temperature_grid[t_idx_max]
        temp_min = self._temperature_grid[t_idx_min]
        fx0 = self._xsec_grid[t_idx_min]
        fx1 = self._xsec_grid[t_idx_max]

        return interp_lin_only(fx0, fx1, temperature, temp_min, temp_max)

    def compute_cia(self, temperature: float) -> npt.NDArray[np.float64]:
        """Computes the collisionally induced absorption cross-section.

        Args:
            temperature: The temperature to compute the opacity at.

        Returns:
            npt.NDArray[np.float64]: The opacity.

        """
        indicies = self.find_closest_temperature_index(temperature)
        return self.interp_linear_grid(temperature, *indicies)

    @property
    def moleculeName(self) -> str:
        """Name of molecule."""
        return self._molecule_name

    @property
    def xsecGrid(self) -> npt.NDArray[np.float64]:
        """Opacity grid."""
        return self._xsec_grid

    @property
    def wavenumberGrid(self) -> npt.NDArray[np.float64]:
        """Wavenumber grid."""
        return self._wavenumber_grid

    @property
    def temperatureGrid(self) -> npt.NDArray[np.float64]:
        """Temperature grid."""
        return self._temperature_grid

moleculeName property

Name of molecule.

temperatureGrid property

Temperature grid.

wavenumberGrid property

Wavenumber grid.

xsecGrid property

Opacity grid.

__init__(molecule_pair, num_t=27, wn_res=15000, wn_size=(300, 30000))

Create the fake opacity.

Parameters:

Name Type Description Default
molecule_pair tuple[str, str]

The pair of molecules.

required
num_t int

The number of temperature points.

27
wn_res int

The resolution of the wavenumber grid.

15000
wn_size tuple[float, float]

The size of the wavenumber grid.

(300, 30000)
Source code in src\taurex_cupy\util.py
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def __init__(
    self,
    molecule_pair: tuple[str, str],
    num_t: int = 27,
    wn_res: int = 15000,
    wn_size: tuple[float, float] = (300, 30000),
) -> None:
    """Create the fake opacity.

    Args:
        molecule_pair: The pair of molecules.
        num_t: The number of temperature points.
        wn_res: The resolution of the wavenumber grid.
        wn_size: The size of the wavenumber grid.



    """
    super().__init__("FAKE", "-".join(molecule_pair))
    self.pair = molecule_pair
    self._wavenumber_grid = create_grid_res(wn_res, *wn_size)[:, 0]
    self._temperature_grid = np.linspace(100, 10000, num_t)
    self._xsec_grid = np.random.rand(self._temperature_grid.size, self._wavenumber_grid.size)

compute_cia(temperature)

Computes the collisionally induced absorption cross-section.

Parameters:

Name Type Description Default
temperature float

The temperature to compute the opacity at.

required

Returns:

Type Description
NDArray[float64]

npt.NDArray[np.float64]: The opacity.

Source code in src\taurex_cupy\util.py
127
128
129
130
131
132
133
134
135
136
137
138
def compute_cia(self, temperature: float) -> npt.NDArray[np.float64]:
    """Computes the collisionally induced absorption cross-section.

    Args:
        temperature: The temperature to compute the opacity at.

    Returns:
        npt.NDArray[np.float64]: The opacity.

    """
    indicies = self.find_closest_temperature_index(temperature)
    return self.interp_linear_grid(temperature, *indicies)

find_closest_temperature_index(temperature)

Finds the nearest indices for a particular temperature

Parameters:

Name Type Description Default
temperature float

The temperature to search for.

required

Returns:

Type Description
tuple[int, int]

tuple[int, int]: The indices of the closest temperatures.

Source code in src\taurex_cupy\util.py
87
88
89
90
91
92
93
94
95
96
97
98
99
def find_closest_temperature_index(self, temperature: float) -> tuple[int, int]:
    """Finds the nearest indices for a particular temperature

    Args:
        temperature: The temperature to search for.

    Returns:
        tuple[int, int]: The indices of the closest temperatures.

    """

    t_min, t_max = find_closest_pair(self.temperatureGrid, temperature)
    return t_min, t_max

interp_linear_grid(temperature, t_idx_min, t_idx_max)

Linear interpolate the CIA opacity.

Parameters:

Name Type Description Default
temperature float

The temperature to interpolate.

required
t_idx_min int

The minimum temperature index.

required
t_idx_max int

The maximum temperature index.

required

Returns:

Type Description
NDArray[float64]

The interpolated opacity.

Source code in src\taurex_cupy\util.py
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
def interp_linear_grid(self, temperature: float, t_idx_min: int, t_idx_max: int) -> npt.NDArray[np.float64]:
    """Linear interpolate the CIA opacity.

    Args:
        temperature: The temperature to interpolate.
        t_idx_min: The minimum temperature index.
        t_idx_max: The maximum temperature index.

    Returns:
        The interpolated opacity.


    """

    if temperature > self._temperature_grid.max():
        return self._xsec_grid[-1]
    elif temperature < self._temperature_grid.min():
        return self._xsec_grid[0]

    temp_max = self._temperature_grid[t_idx_max]
    temp_min = self._temperature_grid[t_idx_min]
    fx0 = self._xsec_grid[t_idx_min]
    fx1 = self._xsec_grid[t_idx_max]

    return interp_lin_only(fx0, fx1, temperature, temp_min, temp_max)

cuda_find_closest_pair(arr, values)

Find the closest pair of values in an array

Parameters

arr : cp.ndarray The array to search values : cp.ndarray The values to search for

Returns

cp.ndarray The indices of the closest pair

Source code in src\taurex_cupy\util.py
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
def cuda_find_closest_pair(arr: cp.ndarray, values: cp.ndarray) -> cp.ndarray:
    """
    Find the closest pair of values in an array

    Parameters
    ----------
    arr : cp.ndarray
        The array to search
    values : cp.ndarray
        The values to search for

    Returns
    -------
    cp.ndarray
        The indices of the closest pair
    """
    right = arr.searchsorted(values, side="right")
    right = cp.clip(right, 0, arr.shape[0] - 1)
    left = right - 1
    left = cp.clip(left, 0, arr.shape[0] - 1)

    return left, right

determine_grid_slice(dest_wngrid, src_wngrid)

Determine the grid length of the destination grid.

Parameters:

Name Type Description Default
dest_wngrid NDArray[float64]

The destination wavenumber grid.

required
src_wngrid NDArray[float64]

The source wavenumber grid.

required

Returns:

Name Type Description
slice slice

The slice of the destination grid.

Source code in src\taurex_cupy\util.py
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
def determine_grid_slice(dest_wngrid: npt.NDArray[np.float64], src_wngrid: npt.NDArray[np.float64]) -> slice:
    """Determine the grid length of the destination grid.

    Args:
        dest_wngrid: The destination wavenumber grid.
        src_wngrid: The source wavenumber grid.

    Returns:
        slice: The slice of the destination grid.

    """
    min_grid_idx = 0
    max_grid_idx = None
    min_wn = dest_wngrid.min()
    max_wn = dest_wngrid.max()
    src_min = src_wngrid.min()
    src_max = src_wngrid.max()
    if min_wn > src_min:
        min_grid_idx = max(np.argmax(min_wn < src_wngrid) - 1, 0)
    if max_wn < src_max:
        max_grid_idx = np.argmax(src_wngrid >= max_wn) + 1

    return slice(min_grid_idx, max_grid_idx)