<commit_before>from .._core.base_components import Diagnostic
from .._core.util import replace_none_with_default
from .._core.array import DataArray
from ._berger_solar_insolation import get_solar_parameters, get_orbital_parameters
import xarray as xr


class BergerSolarInsolation(Diagnostic):
    inputs = ('time', 'longitude', 'latitude')
    diagnostics = (
        'solar_insolation',
        'solar_zenith_angle',
        'obliquity',
        'eccentricity',
        'normalized_earth_sun_distance',
    )

    def __init__(self, solar_constant=None):
        self._orbital_parameters = {}
        self._solar_constant = replace_none_with_default(
            'solar_constant', solar_constant)

    def __call__(self, state):
        """
        Gets diagnostics from the passed model state.

        Args:
            state (dict): A model state dictionary.

        Returns:
            diagnostics (dict): A dictionary whose keys are strings indicating
                state quantities and values are the value of those quantities
                at the time of the input state.

        Raises:
            KeyError: If a required quantity is missing from the state.
            InvalidStateException: If state is not a valid input for the
                Prognostic instance.
        """
        lat, lon = xr.broadcast(
            state['latitude'].to_units('degrees_N'),
            state['longitude'].to_units('degrees_E'))
        if 'solar_constant' in state:
            if state['solar_constant'].size > 1:
                raise ValueError(
                    'Solar constant should be a scalar (inside a DataArray)')
            else:
                solar_constant = state['solar_constant'].values.item()
        else:
            solar_constant = self._solar_constant
        solar_insolation, solar_zenith_angle, obliquity, eccentricity, rho = \
            self._driver(state, lat, lon, solar_constant)
        diagnostics = {
            'solar_insolation': DataArray(
                solar_insolation.reshape(lat.shape),
                dims=lat.dims, attrs={'units': 'W m^-2'}
            ),
            'solar_zenith_angle': DataArray(
                solar_zenith_angle.reshape(lat.shape),
                dims=lat.dims, attrs={'units': 'radians'}
            ),
            'obliquity': DataArray(
                obliquity, dims=[], attrs={'units': ''}
            ),
            'eccentricity': DataArray(
                eccentricity, dims=[], attrs={'units': ''}
            ),
            'normalized_earth_sun_distance': DataArray(
                rho, dims=[], attrs={'units': ''}
            ),
        }
        return diagnostics

    def _driver(self, state, lat, lon, solar_constant):
        # We can compute orbital parameters yearly because they change very
        # slowly
        year = state['time'].year
        if year not in self._orbital_parameters:
            self._orbital_parameters[year] = get_orbital_parameters(
                float(year - 1950))
        lambda_m0, eccentricity, omega_tilde, obliquity = (
            self._orbital_parameters[year])
        return get_solar_parameters(
            lambda_m0,
            eccentricity,
            omega_tilde,
            obliquity,
            years_since_vernal_equinox(state['time']),
            fractional_day(state['time']),
            lat.values.flatten(),
            lon.values.flatten(),
            solar_constant,)



def fractional_year(dt):
    """Returns a float indicating the fraction of the current year that has
    passed in the given datetime object."""
    year_start = type(dt)(dt.year, 1, 1)
    year_end = type(dt)(dt.year+1, 1, 1)
    return (dt - year_start).total_seconds() / (year_end - year_start).total_seconds()


def years_since_jan_1_1950_ad(dt):
    return (dt.year - 1950) + fractional_year(dt)

<commit_msg>Removed extra blank line to pass flake8<commit_after>from .._core.base_components import Diagnostic
from .._core.util import replace_none_with_default
from .._core.array import DataArray
from ._berger_solar_insolation import get_solar_parameters, get_orbital_parameters
import xarray as xr


class BergerSolarInsolation(Diagnostic):
    inputs = ('time', 'longitude', 'latitude')
    diagnostics = (
        'solar_insolation',
        'solar_zenith_angle',
        'obliquity',
        'eccentricity',
        'normalized_earth_sun_distance',
    )

    def __init__(self, solar_constant=None):
        self._orbital_parameters = {}
        self._solar_constant = replace_none_with_default(
            'solar_constant', solar_constant)

    def __call__(self, state):
        """
        Gets diagnostics from the passed model state.

        Args:
            state (dict): A model state dictionary.

        Returns:
            diagnostics (dict): A dictionary whose keys are strings indicating
                state quantities and values are the value of those quantities
                at the time of the input state.

        Raises:
            KeyError: If a required quantity is missing from the state.
            InvalidStateException: If state is not a valid input for the
                Prognostic instance.
        """
        lat, lon = xr.broadcast(
            state['latitude'].to_units('degrees_N'),
            state['longitude'].to_units('degrees_E'))
        if 'solar_constant' in state:
            if state['solar_constant'].size > 1:
                raise ValueError(
                    'Solar constant should be a scalar (inside a DataArray)')
            else:
                solar_constant = state['solar_constant'].values.item()
        else:
            solar_constant = self._solar_constant
        solar_insolation, solar_zenith_angle, obliquity, eccentricity, rho = \
            self._driver(state, lat, lon, solar_constant)
        diagnostics = {
            'solar_insolation': DataArray(
                solar_insolation.reshape(lat.shape),
                dims=lat.dims, attrs={'units': 'W m^-2'}
            ),
            'solar_zenith_angle': DataArray(
                solar_zenith_angle.reshape(lat.shape),
                dims=lat.dims, attrs={'units': 'radians'}
            ),
            'obliquity': DataArray(
                obliquity, dims=[], attrs={'units': ''}
            ),
            'eccentricity': DataArray(
                eccentricity, dims=[], attrs={'units': ''}
            ),
            'normalized_earth_sun_distance': DataArray(
                rho, dims=[], attrs={'units': ''}
            ),
        }
        return diagnostics

    def _driver(self, state, lat, lon, solar_constant):
        # We can compute orbital parameters yearly because they change very
        # slowly
        year = state['time'].year
        if year not in self._orbital_parameters:
            self._orbital_parameters[year] = get_orbital_parameters(
                float(year - 1950))
        lambda_m0, eccentricity, omega_tilde, obliquity = (
            self._orbital_parameters[year])
        return get_solar_parameters(
            lambda_m0,
            eccentricity,
            omega_tilde,
            obliquity,
            years_since_vernal_equinox(state['time']),
            fractional_day(state['time']),
            lat.values.flatten(),
            lon.values.flatten(),
            solar_constant,)


def fractional_year(dt):
    """Returns a float indicating the fraction of the current year that has
    passed in the given datetime object."""
    year_start = type(dt)(dt.year, 1, 1)
    year_end = type(dt)(dt.year+1, 1, 1)
    return (dt - year_start).total_seconds() / (year_end - year_start).total_seconds()


def years_since_jan_1_1950_ad(dt):
    return (dt.year - 1950) + fractional_year(dt)


{'fraction_non_alphanumeric': 0.09686168151879117, 'fraction_numerical': 0.00581170089112747, 'mean_word_length': 2.1301535974130963, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1449826', 'n_tokens_mistral': 2225, 'n_tokens_neox': 2130, 'n_words': 526}