Example #1
0
class _ARTS:
    def __init__(self,
                 ws=None,
                 threads=None,
                 nstreams=4,
                 scale_vmr=True,
                 verbosity=0):
        """Initialize a wrapper for an ARTS workspace.

        Parameters:
            ws (pyarts.workspace.Workspace): An ARTS workspace.
            threads (int): Number of threads to use.
                Default is all available threads.
            nstreams (int): Number of viewing angles to base the radiative
                flux calculation on.
            scale_vmr (bool): Control whether dry volume mixing ratios are
                scaled with the water-vapor concentration (default is `False.`)
            verbosity (int): Control the ARTS verbosity from 0 (quiet) to 2.
        """
        from pyarts.workspace import Workspace, arts_agenda

        self.nstreams = nstreams
        self.scale_vmr = scale_vmr

        if ws is None:
            self.ws = Workspace(verbosity=verbosity)

        self.ws.execute_controlfile("general/general.arts")
        self.ws.execute_controlfile("general/continua.arts")
        self.ws.execute_controlfile("general/agendas.arts")
        self.ws.execute_controlfile("general/planet_earth.arts")

        # Agenda settings
        self.ws.Copy(self.ws.abs_xsec_agenda, self.ws.abs_xsec_agenda__noCIA)
        self.ws.Copy(self.ws.iy_main_agenda, self.ws.iy_main_agenda__Emission)
        self.ws.Copy(self.ws.iy_space_agenda,
                     self.ws.iy_space_agenda__CosmicBackground)
        self.ws.Copy(self.ws.iy_surface_agenda,
                     self.ws.iy_surface_agenda__UseSurfaceRtprop)
        self.ws.Copy(
            self.ws.propmat_clearsky_agenda,
            self.ws.propmat_clearsky_agenda__LookUpTable,
        )
        self.ws.Copy(self.ws.ppath_agenda,
                     self.ws.ppath_agenda__FollowSensorLosPath)
        self.ws.Copy(self.ws.ppath_step_agenda,
                     self.ws.ppath_step_agenda__GeometricPath)

        @arts_agenda
        def p_eq_agenda(workspace):
            workspace.water_p_eq_fieldMK05()

        self.ws.Copy(self.ws.water_p_eq_agenda, p_eq_agenda)

        @arts_agenda
        def cloudbox_agenda(workspace):
            workspace.iyInterpCloudboxField()

        self.ws.Copy(self.ws.iy_cloudbox_agenda, cloudbox_agenda)

        # Number of Stokes components to be computed
        self.ws.IndexSet(self.ws.stokes_dim, 1)

        self.ws.jacobianOff()  # No jacobian calculation
        self.ws.cloudboxOff()  # Clearsky = No scattering

        # Set Absorption Species
        self.ws.abs_speciesSet(species=[
            "O2, O2-CIAfunCKDMT100",
            "H2O, H2O-SelfContCKDMT252, H2O-ForeignContCKDMT252",
            "O3",
            "CO2, CO2-CKDMT252",
            "N2, N2-CIAfunCKDMT252, N2-CIArotCKDMT252",
            "N2O",
            "CH4",
            "CO",
        ])

        # Surface handling
        self.ws.VectorSetConstant(self.ws.surface_scalar_reflectivity, 1, 0.0)
        self.ws.Copy(
            self.ws.surface_rtprop_agenda,
            self.ws.
            surface_rtprop_agenda__Specular_NoPol_ReflFix_SurfTFromt_surface,
        )

        # Read lookup table
        abs_lookup = os.getenv("KONRAD_LOOKUP_TABLE",
                               join(dirname(__file__), "data/abs_lookup.xml"))

        if not isfile(abs_lookup):
            raise FileNotFoundError(
                "Could not find ARTS absorption lookup table.\n"
                "To perform ARTS calculations you have to download the lookup "
                "table at:\n\n    https://doi.org/10.5281/zenodo.3885410\n\n"
                "Afterwards, use the following environment variable to tell "
                "konrad where to find it:\n\n"
                "    $ export KONRAD_LOOKUP_TABLE='/path/to/abs_lookup.xml'")

        self.ws.ReadXML(self.ws.abs_lookup, abs_lookup)
        self.ws.f_gridFromGasAbsLookup()
        self.ws.abs_lookupAdapt()

        # Sensor settings
        self.ws.sensorOff()  # No sensor properties

        # Atmosphere
        self.ws.AtmosphereSet1D()

        # Set number of OMP threads
        if threads is not None:
            self.ws.SetNumberOfThreads(threads)

    def calc_lookup_table(self, filename=None, fnum=2**15, wavenumber=None):
        """Calculate an absorption lookup table.

        The lookup table is constructed to cover surface temperatures
        between 200 and 400 K, and water vapor mixing ratio up to 40%.

        The frequency grid covers the whole outgoing longwave spectrum
        from 10 to 3,250 cm^-1.

        References:
            An absorption lookup table can be found at
                https://doi.org/10.5281/zenodo.3885410

        Parameters:
            filename (str): (Optional) path to an ARTS XML file
                to store the lookup table.
            fnum (int): Number of frequencies in frequency grid.
                Ignored if `wavenumber` is set.
            wavenumber (ndarray): Wavenumber grid [m-1].
        """
        # Create a frequency grid
        if wavenumber is None:
            wavenumber = np.linspace(10e2, 3_250e2, fnum)
        self.ws.f_grid = ty.physics.wavenumber2frequency(wavenumber)

        # Read line catagloge and create absorption lines.
        self.ws.ReadSplitARTSCAT(
            abs_lines=self.ws.abs_lines,
            abs_species=self.ws.abs_species,
            basename="hitran_split_artscat5/",
            fmin=0.0,
            fmax=1e99,
            globalquantumnumbers="",
            localquantumnumbers="",
            ignore_missing=0,
        )

        # Set line shape and cut off.
        self.ws.abs_linesSetLineShapeType(self.ws.abs_lines, "VP")
        self.ws.abs_linesSetNormalization(self.ws.abs_lines, "VVH")
        self.ws.abs_linesSetCutoff(self.ws.abs_lines, "ByLine", 750e9)

        self.ws.abs_lines_per_speciesCreateFromLines()
        self.ws.abs_lines_per_speciesCompact()

        # Create a standard atmosphere
        p_grid = get_quadratic_pgrid(1_200e2, 0.5, 80)

        atmosphere = Atmosphere(p_grid)
        atmosphere["T"][
            -1, :] = 300.0 + 5.0 * np.log(atmosphere["plev"] / 1000e2)
        atmosphere.tracegases_rcemip()
        atmosphere["O2"][:] = 0.2095
        atmosphere["CO2"][:] = 1.5 * 348e-6

        h2o = 0.01 * (p_grid / 1000e2)**0.2
        atmosphere["H2O"][:] = h2o[:-1]

        # Convert the konrad atmosphere into an ARTS atm_fields_compact.
        atm_fields_compact = atmosphere.to_atm_fields_compact()
        self.ws.atm_fields_compact = atm_fields_compact

        self.ws.atm_fields_compactAddConstant(
            atm_fields_compact=self.ws.atm_fields_compact,
            name="abs_species-N2",
            value=0.7808,
            condensibles=["abs_species-H2O"],
        )

        # Setup the lookup table calculation
        self.ws.AtmFieldsAndParticleBulkPropFieldFromCompact()
        self.ws.vmr_field.value = self.ws.vmr_field.value.clip(min=0.0)
        self.ws.atmfields_checkedCalc()
        self.ws.abs_lookupSetup(p_step=1.0)  # Do not refine p_grid
        self.ws.abs_t_pert = np.arange(-160, 61, 20)

        nls_idx = [
            i for i, tag in enumerate(self.ws.abs_species.value)
            if "H2O" in tag[0]
        ]
        self.ws.abs_speciesSet(
            abs_species=self.ws.abs_nls,
            species=[", ".join(self.ws.abs_species.value[nls_idx[0]])],
        )

        self.ws.abs_nls_pert = np.array(
            [10**x for x in [-9, -7, -5, -3, -1, 0, 0.5, 1, 1.5, 2]])

        # Run checks
        self.ws.abs_xsec_agenda_checkedCalc()
        self.ws.lbl_checkedCalc()

        # Calculate actual lookup table.
        self.ws.abs_lookupCalc()

        if filename is not None:
            self.ws.WriteXML("binary", self.ws.abs_lookup, filename)

    def set_atmospheric_state(self, atmosphere, t_surface):
        """Set and check the atmospheric fields."""
        import pyarts

        atm_fields_compact = atmosphere.to_atm_fields_compact()

        # Scale dry-air VMRs with H2O and CO2 content.
        if self.scale_vmr:
            variable_vmrs = (atm_fields_compact.get("abs_species-H2O")[0] +
                             atm_fields_compact.get("abs_species-CO2")[0])
        else:
            t3_shape = atm_fields_compact.get("abs_species-H2O")[0].shape
            variable_vmrs = np.zeros(t3_shape)

        for species in atm_fields_compact.grids[0]:
            if (species.startswith("abs_species-") and "H2O" not in species
                    and "CO2" not in species):
                atm_fields_compact.scale(species, 1 - variable_vmrs)

        # Compute the N2 VMR as a residual of the full atmosphere composition.
        n2 = pyarts.types.GriddedField3(
            grids=atm_fields_compact.grids[1:],
            data=0.7808 * (1 - variable_vmrs),
        )

        self.ws.atm_fields_compact = atm_fields_compact
        self.ws.atm_fields_compactAddSpecies(
            atm_fields_compact=self.ws.atm_fields_compact,
            name="abs_species-N2",
            value=n2,
        )
        self.ws.AtmFieldsAndParticleBulkPropFieldFromCompact()
        self.ws.vmr_field = self.ws.vmr_field.value.clip(min=0)

        # Surface & TOA
        # Add pressure layers to the surface and top-of-the-atmosphere to
        # ensure consistent atmosphere boundaries between ARTS and RRTMG.
        self.ws.t_surface = np.array([[t_surface]])
        self.ws.z_surface = np.array([[0.0]])
        self.ws.z_field.value[0, 0, 0] = 0.0

        # Perform configuration and atmosphere checks
        self.ws.atmfields_checkedCalc()
        self.ws.propmat_clearsky_agenda_checkedCalc()
        self.ws.atmgeom_checkedCalc()
        self.ws.cloudbox_checkedCalc()

    def calc_spectral_irradiance_field(self, atmosphere, t_surface):
        """Calculate the spectral irradiance field."""
        self.set_atmospheric_state(atmosphere, t_surface)

        # get the zenith angle grid and the integrations weights
        self.ws.AngularGridsSetFluxCalc(N_za_grid=self.nstreams,
                                        N_aa_grid=1,
                                        za_grid_type="double_gauss")

        # calculate intensity field
        self.ws.Tensor3Create("trans_field")
        self.ws.spectral_radiance_fieldClearskyPlaneParallel(
            trans_field=self.ws.trans_field,
            use_parallel_za=0,
        )
        self.ws.spectral_irradiance_fieldFromSpectralRadianceField()

        return (
            self.ws.f_grid.value.copy(),
            self.ws.p_grid.value.copy(),
            self.ws.spectral_irradiance_field.value.copy(),
            self.ws.trans_field.value[:, 1:, 0].copy().prod(axis=1),
        )

    def calc_optical_thickness(self, atmosphere, t_surface):
        """Calculate the spectral irradiance field."""
        self.set_atmospheric_state(atmosphere, t_surface)

        self.ws.propmat_clearsky_fieldCalc()

        tau = np.trapz(
            y=self.ws.propmat_clearsky_field.value[:, :, 0, 0, :, 0, 0],
            x=self.ws.z_field.value[:, 0, 0],
            axis=-1,
        )

        return self.ws.f_grid.value.copy(), tau

    @staticmethod
    def integrate_spectral_irradiance(frequency, irradiance):
        """Integrate the spectral irradiance field over the frequency.

        Parameters:
            frequency (ndarray): Frequency [Hz].
            irradiance (ndarray): Spectral irradiance [W m^-2 / Hz].

        Returns:
            ndarray, ndarray: Downward flux, upward, flux [W m^-2]
        """
        F = np.trapz(irradiance, frequency, axis=0)[:, 0, 0, :]

        # Fluxes
        lw_down = -F[:, 0]
        lw_up = F[:, 1]

        return lw_down, lw_up

    def calc_spectral_olr(self, atmosphere, surface):
        """Calculate the outgoing longwave radiation as function of wavenumber.

        Parameters:
            atmosphere (konrad.atmosphere.Atmosphere): Atmosphere model.
            surface (konrad.surface.Surface): Surface model.

        Returns:
           ndarray: Outgoing longwave radiation [W m^-2 / cm^-1]
        """
        f, _, irradiance_field, _ = self.calc_spectral_irradiance_field(
            atmosphere=atmosphere, t_surface=surface["temperature"][0])
        return f, irradiance_field[:, -1, 0, 0, 1]
Example #2
0
class TestVariables:
    """
    Tests the manipulation of workspace variables.
    """
    def setup_method(self):
        """
        This ensures a new Workspace for every test.
        """
        self.dir = os.path.dirname(os.path.realpath(__file__))
        self.ws = Workspace(verbosity=0)
        self.setup_workspace()

    def setup_workspace(self):
        ws = self.ws
        ws.atmosphere_dim = 1
        ws.p_grid = np.linspace(1e5, 1e3, 21)
        ws.Touch(ws.lat_grid)
        ws.Touch(ws.lon_grid)

        ws.f_grid = 183.0e9 * np.ones(1)
        ws.stokes_dim = 1

        ws.sensor_los = 180.0 * np.ones((1, 1))
        ws.sensor_pos = 830e3 * np.ones((1, 1))
        ws.sensorOff()

    def test_index_transfer(self):
        """
        Create and set Index WSV.
        """
        self.ws.IndexCreate("index_variable")
        i = np.random.randint(0, 100)
        self.ws.index_variable = i
        assert self.ws.index_variable.value == i

    def test_string_transfer(self):
        """
        Create and set String WSV.
        """
        self.ws.StringCreate("string_variable")
        s = "some random string."
        self.ws.string_variable = s
        assert self.ws.string_variable.value == s

    def test_array_of_index_transfer(self):
        """
        Create and set ArrayOfIndex WSV.
        """
        self.ws.ArrayOfIndexCreate("array_of_index_variable")
        i = [np.random.randint(0, 100) for j in range(10)]

        self.ws.array_of_index_variable = i
        assert self.ws.array_of_index_variable.value == i

        self.ws.array_of_index_variable = []
        assert self.ws.array_of_index_variable.value == []

    def test_array_of_vector_transfer(self):
        """
        Create and set ArrayOfVector WSV.

        """
        self.ws.ArrayOfVectorCreate("array_of_vector_variable")
        aov = pyarts.xml.load(
            os.path.join(self.dir, "../xml/reference/arrayofvector.xml"))
        self.ws.array_of_vector_variable = aov
        assert self.ws.array_of_vector_variable.value == aov

    def test_vector_transfer(self):
        """
        Create and set Vector WSV.
        """
        self.ws.VectorCreate("vector_variable")
        v = np.random.rand(10)
        self.ws.vector_variable = v
        assert all(self.ws.vector_variable.value == v)

    def test_matrix_transfer(self):
        """
        Create and set Matrix WSV.
        """
        self.ws.MatrixCreate("matrix_variable")
        m = np.random.rand(10, 10)
        self.ws.matrix_variable = m
        assert all(self.ws.matrix_variable.value.ravel() == m.ravel())

    def test_sparse_transfer(self):
        """
        Create and set Sparse WSV.
        """
        n = 100
        d2 = np.ones(n - 2)
        d1 = np.ones(n - 1)
        d = np.ones(n)
        m = sp.sparse.diags(diagonals=[d2, d1, d, d1, d2],
                            offsets=[2, 1, 0, -1, -2])
        self.ws.sensor_response = m
        assert np.all(m.toarray() == self.ws.sensor_response.value.toarray())

    def test_tensor_3(self):
        """
        Create and set Tensor3 variable.
        """
        t_0 = np.random.rand(*([3] * 3))
        self.ws.Tensor3Create("tensor_3")
        self.ws.tensor_3 = t_0
        assert np.all(t_0 == self.ws.tensor_3.value)

    def test_tensor_4(self):
        """
        Create and set Tensor4 variable.
        """
        t_0 = np.random.rand(*([3] * 4))
        t_1 = self.ws.Tensor4Create("tensor_4")
        self.ws.tensor_4 = t_0
        assert np.all(t_0 == self.ws.tensor_4.value)

    def test_tensor_5(self):
        """
        Create and set Tensor5 variable.
        """
        t_0 = np.random.rand(*([3] * 5))
        t_1 = self.ws.Tensor5Create("tensor_5")
        self.ws.tensor_5 = t_0
        assert np.all(t_0 == self.ws.tensor_5.value)

    def test_tensor_6(self):
        """
        Create and set Tensor6 variable.
        """
        t_0 = np.random.rand(*([3] * 6))
        t_1 = self.ws.Tensor6Create("tensor_6")
        self.ws.tensor_6 = t_0
        assert np.all(t_0 == self.ws.tensor_6.value)

    def test_tensor_7(self):
        """
        Create and set Tensor7 variable.
        """
        t_0 = np.random.rand(*([3] * 7))
        self.ws.Tensor7Create("tensor_7")
        self.ws.tensor_7 = t_0
        assert np.all(t_0 == self.ws.tensor_7.value)

    def test_time(self):
        """
        Create and set Time variable.
        """
        times = ["2020-01-02 03:04:05", "2021-02-03 04:05:06"]
        self.ws.ArrayOfTimeCreate("time_1")
        self.ws.ArrayOfTimeNLinSpace(self.ws.time_1, 2, times[0], times[1])
        assert (times[0] == str(self.ws.time_1.value[0])[0:19]
                and times[1] == str(self.ws.time_1.value[1])[0:19])

    def test_creation(self):
        """
        Test creation of WSVs.
        """
        self.ws.ArrayOfIndexCreate("array_of_index")
        self.ws.ArrayOfIndexCreate("array_of_index")
        with pytest.raises(Exception):
            self.ws.VectorCreate("array_of_index")

    def test_covariance_matrix(self):
        """
        Test manipulation of CorvarianceMatrix objects.
        """
        ws = self.ws

        ws.jacobianInit()
        ws.jacobianAddAbsSpecies(species="O3",
                                 g1=ws.p_grid,
                                 g2=ws.lat_grid,
                                 g3=ws.lon_grid)
        ws.jacobianAddAbsSpecies(species="H2O",
                                 g1=ws.p_grid,
                                 g2=ws.lat_grid,
                                 g3=ws.lon_grid)
        ws.jacobianClose()

        ws.covmatDiagonal(out=ws.covmat_block,
                          out_inverse=ws.covmat_block,
                          vars=10.0 * np.ones(ws.p_grid.value.size))
        ws.covmat_sxAddBlock(block=ws.covmat_block)
        ws.covmatDiagonal(out=ws.covmat_block,
                          out_inverse=ws.covmat_block,
                          vars=20.0 * np.ones(ws.p_grid.value.size))
        ws.covmat_sxAddBlock(block=ws.covmat_block)

    def test_variable_creation(self):
        """
        Test creation of named and unnambed WSVs.
        """

        # Unnamed variable
        wsv = self.ws.create_variable("Matrix", None)
        self.ws.__setattr__(wsv.name, np.eye(5))
        assert np.all(
            np.isclose(np.eye(5),
                       self.ws.__getattr__(wsv.name).value))

        # Named variable
        wsv = self.ws.create_variable("Matrix", "matrix_wsv")
        self.ws.matrix_wsv = np.eye(5)
        assert np.all(np.isclose(np.eye(5), self.ws.matrix_wsv.value))

    def test_variable_set_empty(self):
        """
        Test initialization of workspace variables.
        """
        self.ws.f_grid = np.array([94e9])
        self.ws.f_grid = []
        assert self.ws.f_grid.value.size == 0

    def test_variable_create(self):
        """
        Test initialization of workspace variables.
        """
        self.ws = Workspace()
        self.ws.IndexCreate("myindex")
        with pytest.raises(Exception):
            print(self.ws.myindex.value)

    def test_convert(self):
        """
        Test automatic conversion of Python types.
        """

        v = WorkspaceVariable.convert("Index", 1.2)
        assert (v == 1)

        v = WorkspaceVariable.convert("String", "string")
        assert (v == "string")

        v = WorkspaceVariable.convert("Numeric", 1)
        assert (type(v) == np.float64)

        v = WorkspaceVariable.convert("Vector", 1.0)
        assert (v.shape == (1, ))

        v = WorkspaceVariable.convert("Matrix", 1.0)
        assert (v.shape == (1, 1))

        v = WorkspaceVariable.convert("Tensor3", 1.0)
        assert (v.shape == (1, 1, 1))

        v = WorkspaceVariable.convert("Tensor6", 1.0)
        assert (v.shape == (1, 1, 1, 1, 1, 1))

        v = WorkspaceVariable.convert("ArrayOfArrayOfIndex", 1.0)
        assert (type(v) == list)
        assert (type(v[0]) == list)
        assert (type(v[0][0]) == int)

        v = WorkspaceVariable.convert("ArrayOfArrayOfIndex", 1)
        return v