Skip to content

Commit

Permalink
Refactored the integral function for significant speed up by taking a…
Browse files Browse the repository at this point in the history 
…dvantage of that we are no longer using component specific line of sights.
  • Loading branch information
@MetinSa
MetinSa committed Jul 4, 2022
1 parent 15bb431 commit e1edd34
Show file tree
Hide file tree
Showing 3 changed files with 110 additions and 122 deletions.
42 changes: 23 additions & 19 deletions zodipy/_model.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,43 +36,47 @@ class Model:
def n_comps(self) -> int:
return len(self.comps)

def interpolate_source_parameters(
self, comp_label: ComponentLabel, freq: u.Quantity[u.GHz] | u.Quantity[u.m]
) -> tuple[float, float, tuple[float, ...]]:
def interp_source_params(
self, freq: u.Quantity[u.GHz] | u.Quantity[u.m]
) -> tuple[list[float], list[float], tuple[float, ...]]:
"""
Returns interpolated/extrapolated source parameters for a component
given a frequency.
"""
albedos: list[float] = []
emissivities: list[float] = []

freq = freq.to(self.spectrum.unit, equivalencies=u.spectral())
emissivity_interpolator = interp1d(
x=self.spectrum,
y=self.emissivities[comp_label],
fill_value="extrapolate",
)
emissivity = emissivity_interpolator(freq)

if self.albedos is not None:
albedo_interpolator = interp1d(
for comp_label in self.comps.keys():
emissivity_interpolator = interp1d(
x=self.spectrum,
y=self.albedos[comp_label],
y=self.emissivities[comp_label],
fill_value="extrapolate",
)
albedo = albedo_interpolator(freq)
else:
albedo = 0.0
emissivities.append(emissivity_interpolator(freq))

if self.albedos is not None:
albedo_interpolator = interp1d(
x=self.spectrum,
y=self.albedos[comp_label],
fill_value="extrapolate",
)
albedo = albedo_interpolator(freq)
else:
albedo = 0.0
albedos.append(albedo)

if self.phase_coefficients is not None:
phase_coefficient_interpolator = interp1d(
x=self.spectrum,
y=np.asarray(self.phase_coefficients),
fill_value="extrapolate",
)
phase_coefficient = phase_coefficient_interpolator(freq)
phase_coefficients = phase_coefficient_interpolator(freq)
else:
phase_coefficient = [0.0 for _ in range(3)]
phase_coefficients = [0.0 for _ in range(3)]

return emissivity, albedo, tuple(phase_coefficient)
return emissivities, albedos, tuple(phase_coefficients)

def __repr__(self) -> str:
comp_names = [comp_label.value for comp_label in self.comps]
Expand Down
11 changes: 5 additions & 6 deletions zodipy/_source_functions.py
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from __future__ import annotations

from functools import lru_cache
from typing import TypeVar

import astropy.constants as const
import astropy.units as u
Expand All @@ -14,10 +15,10 @@

SPECIFIC_INTENSITY_UNITS = u.W / u.Hz / u.m**2 / u.sr

A = TypeVar("A", float, NDArray[np.floating])

def get_blackbody_emission(
freq: float | NDArray[np.floating], T: float | NDArray[np.floating]
) -> float | NDArray[np.floating]:

def get_blackbody_emission(freq: float, T: A) -> A:
"""Returns the blackbody emission given a frequency.
Parameters
Expand All @@ -39,9 +40,7 @@ def get_blackbody_emission(
return term1 / term2


def get_dust_grain_temperature(
R: float | NDArray[np.floating], T_0: float, delta: float
) -> float | NDArray[np.floating]:
def get_dust_grain_temperature(R: A, T_0: float, delta: float) -> A:
"""Returns the dust grain temperature given a radial distance from the Sun.
Parameters
Expand Down
179 changes: 82 additions & 97 deletions zodipy/zodipy.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,6 @@ class Zodipy:
file. For more information on available ephemeridis, please visit
https://docs.astropy.org/en/stable/coordinates/solarsystem.html
"""

def __init__(
Expand Down Expand Up @@ -386,7 +385,9 @@ def _compute_emission(
if self.model.solar_irradiance_model is not None:
solar_irradiance = (
self.model.solar_irradiance_model.interpolate_solar_irradiance(
freq=freq, albedos=self.model.albedos, extrapolate=self.extrapolate
freq=freq,
albedos=self.model.albedos,
extrapolate=self.extrapolate,
)
)
else:
Expand All @@ -406,6 +407,9 @@ def _compute_emission(
observer_position_expanded = np.expand_dims(observer_position, axis=-1)
earth_position_expanded = np.expand_dims(earth_position, axis=-1)

emissivities, albedos, phase_coefficients = self.model.interp_source_params(
freq
)
# Distribute pointing to available CPUs and compute the emission in parallel.
if self.parallel:
n_proc = multiprocessing.cpu_count() if self.n_proc is None else self.n_proc
Expand All @@ -414,102 +418,77 @@ def _compute_emission(
stop_chunks = np.array_split(stop, n_proc, axis=-1)

with multiprocessing.Pool(processes=n_proc) as pool:
# Compute the integrated line-of-sight emission for each component and
# pointing.
for idx, (label, comp) in enumerate(self.model.comps.items()):
source_parameters = self.model.interpolate_source_parameters(
label, freq
)
emissivity, albedo, phase_coefficients = source_parameters

# Here we create a partial function that will be passed to the
# integration scheme. The arrays are reshaped to (d, n, p) where d
# is the geometrical dimensionality, n is the number of different
# pointings, and p is the number of integration points of the
# quadrature.
partial_integrand_chunks = [
partial(
_compute_comp_emission_at_step,
start=start,
stop=np.expand_dims(stop, axis=-1),
X_obs=observer_position_expanded,
X_earth=earth_position_expanded,
u_los=np.expand_dims(unit_vectors, axis=-1),
comp=comp,
freq=freq.value,
T_0=self.model.T_0,
delta=self.model.delta,
emissivity=emissivity,
albedo=albedo,
phase_coefficients=phase_coefficients,
solar_irradiance=solar_irradiance,
)
for unit_vectors, stop in zip(unit_vector_chunks, stop_chunks)
]

proc_chunks = [
pool.apply_async(
self.integration_scheme.integrate,
args=(partial_integrand, [-1, 1]),
)
for partial_integrand in partial_integrand_chunks
]
integrated_comp_emission = np.concatenate(
[result.get() for result in proc_chunks], axis=0
# Here we create a partial function that will be passed to the
# integration scheme. The arrays are reshaped to (d, n, p) where d
# is the geometrical dimensionality, n is the number of different
# pointings, and p is the number of integration points of the
# quadrature.
partial_integrand_chunks = [
partial(
_compute_emission_at_step,
start=start,
stop=np.expand_dims(stop, axis=-1),
n_quad_points=self.gauss_quad_order,
X_obs=observer_position_expanded,
X_earth=earth_position_expanded,
u_los=np.expand_dims(unit_vectors, axis=-1),
comps=list(self.model.comps.values()),
freq=freq.value,
T_0=self.model.T_0,
delta=self.model.delta,
emissivities=emissivities,
albedos=albedos,
phase_coefficients=phase_coefficients,
solar_irradiance=solar_irradiance,
)
for unit_vectors, stop in zip(unit_vector_chunks, stop_chunks)
]

# Convert the integral from [-1, 1] back to [start, stop].
integrated_comp_emission *= 0.5 * (stop - start)

if binned:
emission[idx, pixels] = integrated_comp_emission
else:
emission[idx] = integrated_comp_emission[indicies]
proc_chunks = [
pool.apply_async(
self.integration_scheme.integrate,
args=(partial_integrand, [-1, 1]),
)
for partial_integrand in partial_integrand_chunks
]
integrated_comp_emission = np.concatenate(
[result.get() for result in proc_chunks], axis=1
)

# Compute the emission in the main process.
# Compute the emission only in the main process.
else:
unit_vectors_expanded = np.expand_dims(unit_vectors, axis=-1)
stop_expanded = np.expand_dims(stop, axis=-1)

# For each component, compute the integrated line of sight emission
for idx, (label, comp) in enumerate(self.model.comps.items()):
source_parameters = self.model.interpolate_source_parameters(
label, freq
)
emissivity, albedo, phase_coefficients = source_parameters
partial_emission_integrand = partial(
_compute_emission_at_step,
start=start,
stop=stop_expanded,
n_quad_points=self.gauss_quad_order,
X_obs=observer_position_expanded,
X_earth=earth_position_expanded,
u_los=unit_vectors_expanded,
comps=list(self.model.comps.values()),
freq=freq.value,
T_0=self.model.T_0,
delta=self.model.delta,
emissivities=emissivities,
albedos=albedos,
phase_coefficients=phase_coefficients,
solar_irradiance=solar_irradiance,
)

# Here we create a partial function that will be passed to the
# integration scheme. The arrays are reshaped to (d, n, p) where d
# is the geometrical dimensionality, n is the number of different
# pointings, and p is the number of integration points of the
# quadrature.
partial_emission_integrand = partial(
_compute_comp_emission_at_step,
start=start,
stop=stop_expanded,
X_obs=observer_position_expanded,
X_earth=earth_position_expanded,
u_los=unit_vectors_expanded,
comp=comp,
freq=freq.value,
T_0=self.model.T_0,
delta=self.model.delta,
emissivity=emissivity,
albedo=albedo,
phase_coefficients=phase_coefficients,
solar_irradiance=solar_irradiance,
)
integrated_comp_emission = self.integration_scheme.integrate(
partial_emission_integrand, [-1, 1]
)

integrated_comp_emission = self.integration_scheme.integrate(
partial_emission_integrand, [-1, 1]
)
# We convert the integral from [-1, 1] back to [start, stop].
integrated_comp_emission *= 0.5 * (stop - start)
# We convert the integral from [-1, 1] back to [start, stop].
integrated_comp_emission *= 0.5 * (stop - start)

if binned:
emission[idx, pixels] = integrated_comp_emission
else:
emission[idx] = integrated_comp_emission[indicies]
if binned:
emission[:, pixels] = integrated_comp_emission
else:
emission = integrated_comp_emission[:, indicies]

if self.solar_cut is not None:
# The observer position is aquired in geocentric coordinates before being
Expand Down Expand Up @@ -538,24 +517,25 @@ def __str__(self) -> str:
return repr(self.model)


def _compute_comp_emission_at_step(
def _compute_emission_at_step(
r: float | NDArray[np.floating],
*,
start: float,
stop: float | NDArray[np.floating],
n_quad_points: int,
X_obs: NDArray[np.floating],
X_earth: NDArray[np.floating],
u_los: NDArray[np.floating],
comp: Component,
comps: list[Component],
freq: float,
T_0: float,
delta: float,
emissivity: float,
albedo: float,
emissivities: list[float],
albedos: list[float],
phase_coefficients: tuple[float, float, float],
solar_irradiance: float,
) -> NDArray[np.floating]:
"""Returns the zodiacal emission at a step along a line of sight."""
"""Returns the zodiacal emission at a step along all lines of sight."""

# Convert the line of sight range from [-1, 1] to the true ecliptic positions
R_los = ((stop - start) / 2) * r + (stop + start) / 2
Expand All @@ -564,17 +544,22 @@ def _compute_comp_emission_at_step(
X_helio = X_los + X_obs
R_helio = np.sqrt(X_helio[0] ** 2 + X_helio[1] ** 2 + X_helio[2] ** 2)

density = comp.compute_density(X_helio, X_earth=X_earth)
temperature = get_dust_grain_temperature(R_helio, T_0, delta)
blackbody_emission = get_blackbody_emission(freq, temperature)

emission = (1 - albedo) * (emissivity * blackbody_emission)
emission = np.zeros((len(comps), np.shape(X_helio)[1], n_quad_points))
density = np.zeros((len(comps), np.shape(X_helio)[1], n_quad_points))

for idx, (comp, albedo, emissivity) in enumerate(zip(comps, albedos, emissivities)):
density[idx] = comp.compute_density(X_helio, X_earth=X_earth)
emission[idx] = (1 - albedo) * (emissivity * blackbody_emission)

if albedo > 0:
if any(albedo != 0 for albedo in albedos):
solar_flux = solar_irradiance / R_helio**2
scattering_angle = get_scattering_angle(R_los, R_helio, X_los, X_helio)
phase_function = get_phase_function(scattering_angle, phase_coefficients)

emission += albedo * solar_flux * phase_function
for idx, albedo in enumerate(albedos):
emission[idx] += albedo * solar_flux * phase_function

return emission * density

0 comments on commit e1edd34

Please sign in to comment.