Source code for pycif.plugins.models.iconart.io.inputs.fluxes

import xarray as xr
import numpy as np
import os
import shutil
from logging import info

from .tracers import change_tracers_xml_fluxes
from .tv_scalef_oem import create_oem_tv_scaling_factors

from ......utils import path
from ......utils.hdf5 import _hdf5_lock




def make_fluxes(self, datastore, ddi, ddf, runsubdir, mode):
    """Native to input function for the oem module"""


    # Inputs and outputs
    trids_in = list(datastore.keys())
    trids_out = [("flux", s) for s in self.chemistry.emis_species]

    # Flags to detect the ensemble
    is_ensemble = False
    is_perturbed_comp = False

    # Create the OEM dir
    oem_dir = os.path.join(runsubdir, 'OEM')
    path.init_dir(oem_dir)

    # Create or use an existing OEM gridded emissions file
    emissions_file = os.path.join(oem_dir, 'gridded_emissions.nc')
    with _hdf5_lock:
        ds_emi = xr.open_dataset(emissions_file) \
            if os.path.exists(emissions_file) else xr.Dataset()
    
    # Loop over the species that must be emitted
    for trid_out in trids_out:

        trid_in = trid_out
        emspec_ref = emspec = trid_out[1]

        # If ensemble, check what the datastore contains
        if "__sample#" in emspec:
            is_ensemble = True
            is_perturbed_comp = True
            
            emspec_ref = emspec.split("__sample#")[0]
            sample_id = emspec.split("__sample#")[1]
            if int(sample_id) > 0:
                continue

            if trid_in not in datastore:
                trid_in = ("flux", emspec_ref)
                is_perturbed_comp = False

        if trid_in not in datastore:
            continue

        tracer = datastore[trid_in]
        tracer_data = tracer["data"][ddi]

        # --------------------------------------------------------------------------
        # -- Get the prior data
        # --------------------------------------------------------------------------

        if "spec" not in tracer_data:
            
            # Fluxes are always loaded to account for the temporal scaling factors
            pass

        else:
            da_flux_prior = tracer_data['spec']
                
            # Adjust the dimensions and remove the time dependency
            da_flux_prior = da_flux_prior.rename({'lon': 'cell'})
            da_flux_prior = da_flux_prior.squeeze()

            # Create OEM scaling factors based on the data
            create_oem_tv_scaling_factors(self, ddi, ddf, da_flux_prior, oem_dir, emspec_ref)

            # TODO: careful, the inversion only works with a single ensemble scaling factor for the period
            da_flux_prior = da_flux_prior.mean(dim='time')

            # Get country ids
            if 'country_ids' not in ds_emi:
                ds_emi['country_ids'] = (
                    'cell', np.arange(da_flux_prior.cell.size))
                
            da_flux_post = da_flux_prior
                
        # --------------------------------------------------------------------------
        # -- Ensemble processing
        # --------------------------------------------------------------------------
        if is_ensemble:

            if is_perturbed_comp:
                info(f"The {emspec_ref} emission category is perturbed.")
                info(f"Calculating emission scaling factors for the {emspec_ref} category...")
                list_da = [datastore[t]["data"][ddi]["spec"] for t in trids_in]
                flux_lambdas = xr.concat(list_da, dim="ens")
                flux_lambdas = flux_lambdas / flux_lambdas[0]
                flux_lambdas = flux_lambdas.rename({'lon': 'reg'})
                flux_lambdas = flux_lambdas.squeeze()
                flux_lambdas = flux_lambdas.isel(time=slice(0, -1)).mean(dim='time')
                flux_lambdas = flux_lambdas.fillna(1)
                flux_lambdas = flux_lambdas.values

            else:
                info(f"The {emspec_ref} emission category is NOT perturbed.")
                info(f"Creating fake emission scaling factors (1.0) for the {emspec_ref} category...")
                nsamples = getattr(self.chemistry, "nsamples", 1)
                flux_lambdas = np.ones((nsamples, self.domain.nlon))

            # Convert the array to a DataArray
            flux_lambdas = xr.DataArray(
                flux_lambdas,
                dims=("ens", "reg")
            )

            # Store the scaling factors
            self.flux_lambdas[ddi][emspec_ref] = flux_lambdas

            # Get the posterior emissions
            flux_lambdas_post = flux_lambdas[2]
            da_flux_post = da_flux_prior * flux_lambdas_post.values

        # --------------------------------------------------------------------------
        # -- Dump the gridded emission files
        # --------------------------------------------------------------------------

        # Just change the attributes
        da_flux_prior.attrs['standard_name'] = f'{emspec_ref}'
        da_flux_prior.attrs['long_name'] = f'Fluxes for {emspec_ref} generated by the CIF'
        da_flux_prior.attrs['units'] = 'kg/m2/s'

        # Add the current category (prior and posterior) to the emission file 
        ds_emi[emspec_ref] = da_flux_prior
        if is_ensemble:
            ds_emi[emspec_ref + '_POST'] = da_flux_post

        # Dump the gridded emissions file
        if os.path.exists(emissions_file):
            os.remove(emissions_file)
        with _hdf5_lock:
            ds_emi.to_netcdf(emissions_file)

        info(f"Added {emspec_ref} to {emissions_file}.")

        # --------------------------------------------------------------------------
        # -- Adapt tracers.xml
        # --------------------------------------------------------------------------
        info(f"Modifying tracers.xml for the {emspec_ref} category ({is_ensemble=})...")
        for t in trids_out:
            change_tracers_xml_fluxes(self, t[1], is_ensemble=is_ensemble)

    return