AURORA / model /aurora_insolation.py

Upload folder using huggingface_hub

87b732d verified 2 months ago

3.66 kB

"""Copyright (c) Microsoft Corporation. Licensed under the MIT license.
Approximate solar insolation computation based on orbital mechanics.
Reference: https://brian-rose.github.io/ClimateLaboratoryBook/courseware/insolation.html
"""

from datetime import datetime
from typing import Sequence, Union

import numpy as np

__all__ = ["insolation"]


def insolation(
    dates: Union[Sequence[datetime], np.ndarray],
    lat: np.ndarray,
    lon: np.ndarray,
    s0: float = 1.0,
    daily: bool = False,
    enforce_2d: bool = False,
    clip_zero: bool = False,
) -> np.ndarray:
    """Calculate approximate solar insolation for given dates, latitudes, and longitudes.
    Uses 1995 orbital elements (standard in the climate modelling community).
    Args:
        dates: 1-D sequence of datetime-like objects.
        lat: 1-D or 2-D array of latitudes in degrees (`-90` to `90`).
        lon: 1-D or 2-D array of longitudes in degrees (`0` to `360`). If 2-D, must have the same
            shape as `lat`.
        s0: Scaling factor (solar constant). Defaults to `1.0`.
        daily: If `True`, return the daily maximum solar radiation (depends only on latitude and day
            of year). Defaults to `False`.
        enforce_2d: If `True` and `lat` / `lon` are 1-D, broadcast them into 2-D meshgrids. Defaults
            to `False`.
        clip_zero: If `True`, set negative (night-time) values to zero. Defaults to `False`.
    Returns:
        :class:`numpy.ndarray`: Insolation array of shape `(len(dates), *lat.shape)`.
    """
    if lat.ndim != lon.ndim:
        raise ValueError("`lat` and `lon` must have the same number of dimensions.")
    if lat.ndim == 1 and enforce_2d:
        lon, lat = np.meshgrid(lon, lat)
    if lat.shape != lon.shape:  # Assert same shape after potential broadcasting
        raise ValueError(f"Shape mismatch between `lat` (`{lat.shape}`) and `lon` (`{lon.shape}`).")
    n_dim = len(lat.shape)
    lat = lat.astype(np.float32)  # No mutation - safe to redefine the local variable.

    # Constants for year 1995.
    eps = 23.4441 * np.pi / 180.0  # Obliquity of Earth
    ecc = 0.016715  # Eccentricity of Earth's orbit
    om = 282.7 * np.pi / 180.0  # Longitude of perihelion
    beta = np.sqrt(1 - ecc**2.0)

    # Day of year as a float.
    dates_arr = np.array(dates, dtype="datetime64")
    start_years = dates_arr.astype("datetime64[Y]")
    days_arr = ((dates_arr - start_years) / np.timedelta64(1, "D")).astype(np.float32)
    for _ in range(n_dim):
        days_arr = np.expand_dims(days_arr, -1)

    if daily:
        days_arr = 0.5 + np.round(days_arr)
        new_lon = lon.astype(np.float32, copy=True)  # Copy to safely mutate.
        new_lon[:] = 0.0
    else:
        new_lon = lon.astype(np.float32)  # No mutation - safe to redefine the local variable.

    # Longitude of Earth relative to the orbit (1st-order approximation).
    lambda_m0 = ecc * (1.0 + beta) * np.sin(om)
    lambda_m = lambda_m0 + 2.0 * np.pi * (days_arr - 80.5) / 365.0
    lambda_ = lambda_m + 2.0 * ecc * np.sin(lambda_m - om)

    # Solar declination.
    dec = np.arcsin(np.sin(eps) * np.sin(lambda_)).astype("float32")
    # Hour angle.
    h = (2 * np.pi * (days_arr + new_lon / 360.0)).astype("float32", copy=False)
    # Earth-Sun distance factor.
    rho = ((1.0 - ecc**2.0) / (1.0 + ecc * np.cos(lambda_ - om))).astype("float32", copy=False)

    # Insolation.
    diff = np.sin(np.pi / 180.0 * lat[None, ...]) * np.sin(dec) - np.cos(
        np.pi / 180.0 * lat[None, ...]
    ) * np.cos(dec) * np.cos(h)
    sol = s0 * diff * rho**-2.0
    if clip_zero:
        sol[sol < 0.0] = 0.0

    return sol

	"""Copyright (c) Microsoft Corporation. Licensed under the MIT license.

	Approximate solar insolation computation based on orbital mechanics.
	Reference: https://brian-rose.github.io/ClimateLaboratoryBook/courseware/insolation.html
	"""

	from datetime import datetime
	from typing import Sequence, Union

	import numpy as np

	__all__ = ["insolation"]


	def insolation(
	dates: Union[Sequence[datetime], np.ndarray],
	lat: np.ndarray,
	lon: np.ndarray,
	s0: float = 1.0,
	daily: bool = False,
	enforce_2d: bool = False,
	clip_zero: bool = False,
	) -> np.ndarray:
	"""Calculate approximate solar insolation for given dates, latitudes, and longitudes.

	Uses 1995 orbital elements (standard in the climate modelling community).

	Args:
	dates: 1-D sequence of datetime-like objects.
	lat: 1-D or 2-D array of latitudes in degrees (`-90` to `90`).
	lon: 1-D or 2-D array of longitudes in degrees (`0` to `360`). If 2-D, must have the same
	shape as `lat`.
	s0: Scaling factor (solar constant). Defaults to `1.0`.
	daily: If `True`, return the daily maximum solar radiation (depends only on latitude and day
	of year). Defaults to `False`.
	enforce_2d: If `True` and `lat` / `lon` are 1-D, broadcast them into 2-D meshgrids. Defaults
	to `False`.
	clip_zero: If `True`, set negative (night-time) values to zero. Defaults to `False`.

	Returns:
	:class:`numpy.ndarray`: Insolation array of shape `(len(dates), *lat.shape)`.
	"""
	if lat.ndim != lon.ndim:
	raise ValueError("`lat` and `lon` must have the same number of dimensions.")
	if lat.ndim == 1 and enforce_2d:
	lon, lat = np.meshgrid(lon, lat)
	if lat.shape != lon.shape: # Assert same shape after potential broadcasting
	raise ValueError(f"Shape mismatch between `lat` (`{lat.shape}`) and `lon` (`{lon.shape}`).")
	n_dim = len(lat.shape)
	lat = lat.astype(np.float32) # No mutation - safe to redefine the local variable.

	# Constants for year 1995.
	eps = 23.4441 * np.pi / 180.0 # Obliquity of Earth
	ecc = 0.016715 # Eccentricity of Earth's orbit
	om = 282.7 * np.pi / 180.0 # Longitude of perihelion
	beta = np.sqrt(1 - ecc**2.0)

	# Day of year as a float.
	dates_arr = np.array(dates, dtype="datetime64")
	start_years = dates_arr.astype("datetime64[Y]")
	days_arr = ((dates_arr - start_years) / np.timedelta64(1, "D")).astype(np.float32)
	for _ in range(n_dim):
	days_arr = np.expand_dims(days_arr, -1)

	if daily:
	days_arr = 0.5 + np.round(days_arr)
	new_lon = lon.astype(np.float32, copy=True) # Copy to safely mutate.
	new_lon[:] = 0.0
	else:
	new_lon = lon.astype(np.float32) # No mutation - safe to redefine the local variable.

	# Longitude of Earth relative to the orbit (1st-order approximation).
	lambda_m0 = ecc * (1.0 + beta) * np.sin(om)
	lambda_m = lambda_m0 + 2.0 * np.pi * (days_arr - 80.5) / 365.0
	lambda_ = lambda_m + 2.0 * ecc * np.sin(lambda_m - om)

	# Solar declination.
	dec = np.arcsin(np.sin(eps) * np.sin(lambda_)).astype("float32")
	# Hour angle.
	h = (2 * np.pi * (days_arr + new_lon / 360.0)).astype("float32", copy=False)
	# Earth-Sun distance factor.
	rho = ((1.0 - ecc*2.0) / (1.0 + ecc np.cos(lambda_ - om))).astype("float32", copy=False)

	# Insolation.
	diff = np.sin(np.pi / 180.0 * lat[None, ...]) * np.sin(dec) - np.cos(
	np.pi / 180.0 * lat[None, ...]
	) * np.cos(dec) * np.cos(h)
	sol = s0 * diff * rho**-2.0
	if clip_zero:
	sol[sol < 0.0] = 0.0

	return sol