Ejemplo n.º 1
0
class ArtsController():
    """A not so high level interface to ARTS."""

    def __init__(self, verbosity=0, agenda_verbosity=0):
        self.ws = Workspace(verbosity=verbosity, agenda_verbosity=agenda_verbosity)
        self.retrieval_quantities = []
        self._sensor = None
        self._observations = []

    def setup(self, atmosphere_dim=1, iy_unit='RJBT', ppath_lmax=-1, stokes_dim=1):
        """
        Run boilerplate (includes, agendas) and set basic variables.

        :param atmosphere_dim:
        :param iy_unit:
        :param ppath_lmax:
        :param stokes_dim:
        """
        boilerplate.include_general(self.ws)
        boilerplate.copy_agendas(self.ws)
        boilerplate.set_basics(self.ws, atmosphere_dim, iy_unit, ppath_lmax, stokes_dim)

        # Deactivate some stuff (can be activated later)
        self.ws.jacobianOff()
        self.ws.cloudboxOff()

    def checked_calc(self, negative_vmr_ok=False, bad_partition_functions_ok=False):
        """Run checked calculations."""
        boilerplate.run_checks(self.ws, negative_vmr_ok, bad_partition_functions_ok)

    def set_spectroscopy(self, abs_lines, abs_species, line_shape=None, abs_f_interp_order=3):
        """
        Setup absorption species and spectroscopy data.

        :param ws: The workspace.
        :param abs_lines: Absoption lines.
        :param abs_species: List of abs species tags.
        :param line_shape: Line shape definition. Default: ['Voigt_Kuntz6', 'VVH', 750e9]
        :param f_abs_interp_order: No effect for OnTheFly propmat. Default: 3
        :type abs_lines: typhon.arts.catalogues.ArrayOfLineRecord
        """
        boilerplate.setup_spectroscopy(self.ws, abs_lines, abs_species, line_shape)
        self.ws.abs_f_interp_order = abs_f_interp_order  # no effect for OnTheFly propmat

    def set_spectroscopy_from_file(self, abs_lines_file, abs_species, format='Arts', line_shape=None, abs_f_interp_order=3):
        """
        Setup absorption species and spectroscopy data from XML file.

        :param ws: The workspace.
        :param abs_lines_file: Path to an XML file.
        :param abs_species: List of abs species tags.
        :param format: One of 'Arts', 'Jpl', 'Hitran' (and others for which a WSM `abs_linesReadFrom...` exists)
        :param line_shape: Line shape definition. Default: ['Voigt_Kuntz6', 'VVH', 750e9]
        :param f_abs_interp_order: No effect for OnTheFly propmat. Default: 3
        """
        ws = self.ws
        if line_shape is None:
            line_shape = ['Voigt_Kuntz6', 'VVH', 750e9]
        ws.abs_speciesSet(abs_species)
        ws.abs_lineshapeDefine(*line_shape)
        #ws.ReadXML(ws.abs_lines, abs_lines_file)
        read_fn = getattr(ws, 'abs_linesReadFrom' + format)
        read_fn(filename=abs_lines_file, fmin=float(0), fmax=float(10e12))
        ws.abs_lines_per_speciesCreateFromLines()
        ws.abs_f_interp_order = abs_f_interp_order

    def set_grids(self, f_grid, p_grid, lat_grid=None, lon_grid=None):
        """
        Set the forward model grids. Basic checks are performed, depending on dimensionality of atmosphere.

        :param f_grid:
        :param p_grid:
        :param lat_grid:
        :param lon_grid:
        """
        if lat_grid is None:
            lat_grid = np.array([0])
        if lon_grid is None:
            lon_grid = np.array([0])

        if not self.ws.atmosphere_dim.initialized:
            raise Exception('atmosphere_dim must be initialized before assigning grids.')
        if not _is_asc(f_grid):
            raise ValueError('Values of f_grid must be strictly increasing.')
        if not _is_desc(p_grid) or not np.all(p_grid > 0):
            raise ValueError('Values of p_grid must be strictly decreasing and positive.')
        if self.ws.atmosphere_dim.value == 1:
            if lat_grid.size > 1 or lon_grid.size > 1:
                raise ValueError('For 1D atmosphere, lat_grid and lon_grid shall be of length 1.')
        elif self.ws.atmosphere_dim.value == 2:
            if lon_grid is not None or len(lat_grid):
                raise ValueError('For 2D atmosphere, lon_grid shall be empty.')
            if lat_grid is None or len(lat_grid) == 0:
                raise ValueError('For 2D atmosphere, lat_grid must be set.')
        elif self.ws.atmosphere_dim.value == 3:
            if lat_grid is None or len(lat_grid) < 2 or lon_grid is None or len(lon_grid) < 2:
                raise ValueError('For 3D atmosphere, length of lat_grid and lon_grid must be >= 2.')
            if max(abs(lon_grid)) > 360:
                raise ValueError('Values of lon_grid must be in the range [-360,360].')
            if max(abs(lat_grid)) > 90:
                raise ValueError('Values of lat_grid must be in the range [-90,90].')
        if lat_grid is not None and not _is_asc(lat_grid):
            raise ValueError('Values of lat_grid must be strictly increasing.')
        if lon_grid is not None and not _is_asc(lon_grid):
            raise ValueError('Values of lon_grid must be strictly increasing.')

        self.ws.f_grid = f_grid
        self.ws.p_grid = p_grid
        if self.atmosphere_dim > 1:
            self.ws.lat_grid = lat_grid
            self.ws.lon_grid = lon_grid
        else:
            self.ws.lat_grid = []
            self.ws.lon_grid = []
            self.ws.lat_true = lat_grid
            self.ws.lon_true = lon_grid
        self.set_surface(0)

    def set_surface(self, altitude: float):
        """
        Set surface altitude.

        :param float altitude:

        .. note:: Currently, only constant altitudes are supported.
        """
        shape = max((1, 1), (self.n_lat, self.n_lat))
        self.ws.z_surface = altitude * np.ones(shape)

    def set_atmosphere(self, atmosphere, vmr_zeropadding=False):
        """
        Set the atmospheric state.

        :param atmosphere: Atmosphere with Temperature, Altitude and VMR Fields.
        :type atmosphere: retrievals.arts.atmosphere.Atmosphere
        :param vmr_zeropadding: Allow VMR zero padding wind fields are always zero
                                padded. Default: False.

        .. note:: Currently only supports 1D atmospheres that is then expanded to a
            multi-dimensional homogeneous atmosphere.
        """
        vmr_zeropadding = 1 if vmr_zeropadding else 0

        self.ws.t_field_raw = atmosphere.t_field
        self.ws.z_field_raw = atmosphere.z_field
        self.ws.vmr_field_raw = [atmosphere.vmr_field(mt) for mt in self.abs_species_maintags]
        self.ws.nlte_field_raw = None

        for c in ('u', 'v', 'w'):
            try:
                raw_field = atmosphere.wind_field(c)
                field = p_interpolate(
                    self.p_grid, raw_field.grids[0], raw_field.data[:, 0, 0], fill=0
                )
            except KeyError:
                field = np.zeros_like(self.p_grid)
            field = np.tile(
                field[:, np.newaxis, np.newaxis], (1, self.n_lat, self.n_lon)
            )
            field_name = 'wind_{}_field'.format(c)
            setattr(self.ws, field_name, field)

        if self.atmosphere_dim == 1:
            self.ws.AtmFieldsCalc(vmr_zeropadding=vmr_zeropadding)
        else:
            self.ws.AtmFieldsCalcExpand1D(vmr_zeropadding=vmr_zeropadding)

    def apply_hse(self, p_hse=100e2, z_hse_accuracy=0.5):
        """
        Calculate z field from hydrostatic equilibrium. See :arts:method:`z_fieldFromHSE`.

        :param p_hse: See :arts:variable:`p_hse`.
        :param z_hse_accuracy: See :arts:variable:`z_hse_accuracy`.
        """
        ws = self.ws
        ws.p_hse = p_hse
        ws.z_hse_accuracy = z_hse_accuracy
        ws.atmfields_checkedCalc()
        ws.z_fieldFromHSE()

    def set_wind(self, wind_u=None, wind_v=None, wind_w=None):
        """
        Set the wind fields to constant values.
        """
        ws = self.ws
        n_p, n_lat, n_lon = self.n_p, self.n_lat, self.n_lon

        if wind_u is not None:
            ws.Tensor3SetConstant(ws.wind_u_field, n_p, n_lat, n_lon, float(wind_u))
        if wind_v is not None:
            ws.Tensor3SetConstant(ws.wind_v_field, n_p, n_lat, n_lon, float(wind_v))
        if wind_w is not None:
            ws.Tensor3SetConstant(ws.wind_w_field, n_p, n_lat, n_lon, float(wind_w))

    def _check_set_wind(self):
        """Set uninitialized wind fields to 0."""
        ws = self.ws
        if ws.wind_u_field.value.size == 0:
            self.set_wind(wind_u=0)
        if ws.wind_v_field.value.size == 0:
            self.set_wind(wind_v=0)
        if ws.wind_w_field.value.size == 0:
            self.set_wind(wind_w=0)

    def set_observations(self, observations):
        """
        Set the geometry of the observations made.

        :param observations:
        :type observations: Iterable[retrievals.arts.interface.Observation]
        """
        sensor_los = np.array([[obs.za, obs.aa] for obs in observations])
        sensor_pos = np.array([[obs.alt, obs.lat, obs.lon] for obs in observations])
        self.ws.sensor_los = sensor_los[:, :max(self.atmosphere_dim - 1, 1)]
        self.ws.sensor_pos = sensor_pos[:, :self.atmosphere_dim]
        self.ws.sensor_time = np.array([obs.time for obs in observations])
        self._observations = observations

    def set_y(self, ys):
        """
        Set the observations.

        :param ys: List with a spectrum for every observation.
        :return:
        """
        y = np.concatenate(ys)
        self.ws.y = y

    def set_sensor(self, sensor):
        """
        Set the sensor.

        :param sensor:
        :type sensor: retrievals.arts.sensors.AbstractSensor
        """
        self._sensor = sensor
        sensor.apply(self.ws)

    def y_calc(self, jacobian_do=False):
        """
        Run the forward model.

        :param jacobian_do: Not implemented yet.
        :return: The measurements as list with length according to observations.
        """
        if jacobian_do:
            raise NotImplementedError('Jacobian not implemented yet.')
            # self.ws.jacobian_do = 1
        self.ws.yCalc()
        return self.y

    def define_retrieval(self, retrieval_quantities, y_vars):
        """
        Define the retrieval quantities.

        :param retrieval_quantities: Iterable of retrieval quantities `retrievals.arts.retrieval.RetrievalQuantity`.
        :param y_vars: List or variance vectors according.
        """
        ws = self.ws

        if isinstance(y_vars, (list, tuple)):
            y_vars = np.concatenate(y_vars)

        if len(y_vars) != self.n_y:
            raise ValueError('Variance vector y_vars must have same length as y.')
        ws.retrievalDefInit()

        # Retrieval quantities
        self.retrieval_quantities = retrieval_quantities
        for rq in retrieval_quantities:
            rq.apply(ws)

        # Se and its inverse
        covmat_block = sparse.diags(y_vars, format='csr')
        boilerplate.set_variable_by_xml(ws, ws.covmat_block, covmat_block)
        ws.covmat_seAddBlock(block=ws.covmat_block)

        covmat_block = sparse.diags(1/y_vars, format='csr')
        boilerplate.set_variable_by_xml(ws, ws.covmat_block, covmat_block)
        ws.covmat_seAddInverseBlock(block=ws.covmat_block)

        ws.retrievalDefClose()

    def oem(self, method='li', max_iter=10, stop_dx=0.01, lm_ga_settings=None, display_progress=True,
            inversion_iterate_agenda=None):
        """
        Run the optimal estimation. See Arts documentation for details.

        :param method:
        :param max_iter:
        :param stop_dx:
        :param lm_ga_settings: Default: [10, 2, 2, 100, 1, 99]
        :param display_progress:
        :param inversion_iterate_agenda: If set to None, a simple default agenda is used.
        :return:
        """
        ws = self.ws

        if lm_ga_settings is None:
            lm_ga_settings = [100.0, 2.0, 2.0, 10.0, 1.0, 1.0]
        lm_ga_settings = np.array(lm_ga_settings)

        # x, jacobian and yf must be initialised
        ws.x = np.array([])
        ws.yf = np.array([])
        ws.jacobian = np.array([[]])

        if inversion_iterate_agenda is None:
            inversion_iterate_agenda = boilerplate.inversion_iterate_agenda
        ws.Copy(ws.inversion_iterate_agenda, inversion_iterate_agenda)

        ws.AgendaExecute(ws.sensor_response_agenda)

        # a priori values
        self._check_set_wind()
        ws.xaStandard()
        xa = ws.xa.value
        for rq in self.retrieval_quantities:
            rq.extract_apriori(xa)

        # Run inversion
        try:
            ws.OEM(method=method, max_iter=max_iter, stop_dx=stop_dx, lm_ga_settings=lm_ga_settings,
                   display_progress=1 if display_progress else 0)
        except Exception as e:
            raise OemException(e)

        if not self.oem_converged:  # Just checks if dxdy is initialized
            return False

        ws.x2artsAtmAndSurf()
        ws.x2artsSensor()
        ws.avkCalc()
        ws.covmat_ssCalc()
        ws.covmat_soCalc()
        ws.retrievalErrorsExtract()

        x = ws.x.value
        avk = ws.avk.value
        eo = ws.retrieval_eo.value
        es = ws.retrieval_ss.value

        for rq in self.retrieval_quantities:
            rq.extract_result(x, avk, eo, es)

        return True

    def get_level2_xarray(self):
        ds = xr.merge([rq.to_xarray() for rq in self.retrieval_quantities])

        # Spectra
        f_backend = self._sensor.f_backend if self._sensor.f_backend is not None else self.f_grid
        ds['f'] = ('f', f_backend)
        ds['y'] = (('observation', 'f'), np.stack(self.y))
        ds['yf'] = (('observation', 'f'), np.stack(self.yf))
        ds['oem_diagnostics'] = ('oem_diagnostics_idx', self.oem_diagnostics)

        y_baseline = self.y_baseline
        if y_baseline is not None:
            ds['y_baseline'] = (('observation', 'f'), np.stack(self.y_baseline))

        # Observations
        for f in Observation._fields:
            values = np.array([getattr(obs, f) for obs in self._observations], dtype=np.float)
            ds['obs_'+f] = ('observation', values)

        # Global attributes
        ds.attrs['arts_version'] = self.arts_version
        ds.attrs['uuid'] = str(uuid.uuid4())
        ds.attrs['date_created'] = datetime.datetime.utcnow().isoformat()
        ds.attrs['nodename'] = os.uname().nodename

        return ds

    @property
    def p_grid(self):
        return self.ws.p_grid.value

    @property
    def lat_grid(self):
        return self.ws.lat_grid.value

    @property
    def lon_grid(self):
        return self.ws.lon_grid.value

    @property
    def f_grid(self):
        return self.ws.f_grid.value

    @property
    def n_p(self):
        return len(self.ws.p_grid.value)

    @property
    def n_lat(self):
        if self.atmosphere_dim == 1:
            return 1
        return len(self.ws.lat_grid.value)

    @property
    def n_lon(self):
        if self.atmosphere_dim == 1:
            return 1
        return len(self.ws.lat_grid.value)

    @property
    def n_y(self):
        return len(self.ws.y.value)

    @property
    def y(self):
        y = np.copy(self.ws.y.value)
        return np.split(y, self.n_obs)

    @property
    def yf(self):
        yf = np.copy(self.ws.yf.value)
        return np.split(yf, self.n_obs)

    @property
    def y_baseline(self):
        if self.ws.y_baseline.initialized:
            bl = np.copy(self.ws.y_baseline.value)
            if bl.size == 1 and bl[0] == 0:
                return None
            return np.split(bl, self.n_obs)
        else:
            return None

    @property
    def n_obs(self):
        if self.ws.sensor_time.initialized:
            return len(self.ws.sensor_time.value)
        else:
            return 0

    @property
    def abs_species_maintags(self):
        abs_species = self.ws.abs_species.value
        maintags = [st[0].split('-')[0] for st in abs_species]
        return maintags

    @property
    def atmosphere_dim(self):
        return self.ws.atmosphere_dim.value

    @property
    def oem_converged(self):
        return self.ws.oem_diagnostics.value[0] == 0

    @property
    def oem_diagnostics(self):
        return self.ws.oem_diagnostics.value

    @property
    def oem_errors(self):
        return self.ws.oem_errors.value

    @property
    def arts_version(self):
        cmd = os.path.join(os.environ['ARTS_BUILD_PATH'], 'src/arts')
        output = os.popen(cmd + ' --version').read().splitlines()
        return output[0]
Ejemplo n.º 2
0
# Atmosphere (A Priori)
# We create a pressure grid using the `PFromZSimple` function to create a grid of approximate pressure levels
# corresponding to altitudes in the range
# z = 0.0, 2000.0, ..., 94000.0
z_toa = 95e3
z_surf = 1e3
z_grid = np.arange(z_surf - 1e3, z_toa, 2e3)
ws.PFromZSimple(ws.p_grid, z_grid)
ws.lat_grid = np.arange(-40.0, 1.0, 40.0)
ws.lon_grid = np.arange(40.0, 61.0, 20.0)
ws.z_surface = z_surf * np.ones(
    (np.asarray(ws.lat_grid).size, np.asarray(ws.lon_grid).size))

# For the a priori state we read data from the Fascod climatology that is part of the ARTS xml data.
ws.AtmRawRead(basename="planets/Earth/Fascod/tropical/tropical")
ws.AtmFieldsCalcExpand1D()

# Adding Wind
# Wind in ARTS is represented by the `wind_u_field` and `wind_v_field` WSVs, which hold the horizontal components
# of the wind at each grid point of the atmosphere model. For this example, a constant wind is assumed.
u_wind = 60.0
v_wind = -40.0
ws.wind_u_field = u_wind * np.ones(
    (ws.p_grid.value.size, ws.lat_grid.value.size, ws.lon_grid.value.size))
ws.wind_v_field = v_wind * np.ones(
    (ws.p_grid.value.size, ws.lat_grid.value.size, ws.lon_grid.value.size))
ws.wind_w_field = np.zeros((0, 0, 0))

# Frequency Grid and Sensor
# The frequency grid for the simulation consists of 119 grid points between 110.516 and 111.156 GHz. The frequencies
# are given by a degree-10 polynomial that has been obtained from a fit to the data from the original `qpack` example.
Ejemplo n.º 3
0
def test_wind_3d_demo():
    ws = Workspace()

    ws.execute_controlfile("general/general.arts")
    ws.verbositySet(0, 0, 0, 0)
    ws.execute_controlfile("general/agendas.arts")
    ws.execute_controlfile("general/continua.arts")
    ws.execute_controlfile("general/planet_earth.arts")

    ws.Copy(ws.abs_xsec_agenda, ws.abs_xsec_agenda__noCIA)
    ws.Copy(ws.ppath_agenda, ws.ppath_agenda__FollowSensorLosPath)
    ws.Copy(ws.ppath_step_agenda, ws.ppath_step_agenda__GeometricPath)
    ws.Copy(ws.iy_space_agenda, ws.iy_space_agenda__CosmicBackground)
    ws.Copy(ws.iy_surface_agenda, ws.iy_surface_agenda__UseSurfaceRtprop)
    ws.Copy(ws.iy_main_agenda, ws.iy_main_agenda__Emission)
    ws.Copy(ws.propmat_clearsky_agenda, ws.propmat_clearsky_agenda__OnTheFly)

    # General Settings
    # For the wind retrievals, the forward model calculations are performed on a 3D atmosphere grid.
    # Radiation is assumed to be unpolarized.
    ws.atmosphere_dim = 3
    ws.stokes_dim = 1
    ws.iy_unit = "RJBT"

    # Absorption
    # We only consider absorption from ozone in this example. The lineshape data is available from
    # the ARTS testdata available in `controlfiles/testdata`.
    ws.abs_speciesSet(["O3", "H2O-PWR98"])
    ws.abs_lineshapeDefine("Voigt_Kuntz6", "VVH", 750e9)
    ws.ReadXML(ws.abs_lines, "testdata/ozone_line.xml")
    ws.abs_lines_per_speciesCreateFromLines()

    # Atmosphere (A Priori)
    # We create a pressure grid using the `PFromZSimple` function to create a grid of approximate pressure levels
    # corresponding to altitudes in the range
    # z = 0.0, 2000.0, ..., 94000.0
    z_toa = 95e3
    z_surf = 1e3
    z_grid = np.arange(z_surf - 1e3, z_toa, 2e3)
    ws.PFromZSimple(ws.p_grid, z_grid)
    ws.lat_grid = np.arange(-40.0, 1.0, 40.0)
    ws.lon_grid = np.arange(40.0, 61.0, 20.0)
    ws.z_surface = z_surf * np.ones(
        (np.asarray(ws.lat_grid).size, np.asarray(ws.lon_grid).size))

    # For the a priori state we read data from the Fascod climatology that is part of the ARTS xml data.
    ws.AtmRawRead(basename="planets/Earth/Fascod/tropical/tropical")
    ws.AtmFieldsCalcExpand1D()

    # Adding Wind
    # Wind in ARTS is represented by the `wind_u_field` and `wind_v_field` WSVs, which hold the horizontal components
    # of the wind at each grid point of the atmosphere model. For this example, a constant wind is assumed.
    u_wind = 60.0
    v_wind = -40.0
    ws.wind_u_field = u_wind * np.ones(
        (ws.p_grid.value.size, ws.lat_grid.value.size, ws.lon_grid.value.size))
    ws.wind_v_field = v_wind * np.ones(
        (ws.p_grid.value.size, ws.lat_grid.value.size, ws.lon_grid.value.size))
    ws.wind_w_field = np.zeros((0, 0, 0))

    # Frequency Grid and Sensor
    # The frequency grid for the simulation consists of 119 grid points between 110.516 and 111.156 GHz.
    # The frequencies are given by a degree-10 polynomial that has been obtained from a fit to the data from
    # the original `qpack` example. This is obscure but also kind of cool.
    coeffs = np.array([
        5.06312189e-08, -2.68851772e-05, 6.20655463e-03, -8.16344090e-01,
        6.75337174e+01, -3.66786505e+03, 1.32578167e+05, -3.14514304e+06,
        4.57491354e+07, 1.10516484e+11
    ])
    ws.f_grid = np.poly1d(coeffs)(np.arange(119))

    # For the sensor we assume a channel width and channel spacing of 50 kHz. We also call AntennaOff to compute
    # only one pencilbeam along the line of sight of the sensor.
    df = 50e3
    f_backend = np.arange(ws.f_grid.value.min() + 2.0 * df,
                          ws.f_grid.value.max() - 2.0 * df, df)
    ws.backend_channel_responseGaussian(np.array([df]), np.array([2.0]))
    ws.AntennaOff()

    ws.sensor_norm = 1
    ws.sensor_time = np.zeros(1)
    ws.sensor_responseInit()

    # Sensor Position and Viewing Geometry
    # 5 Measurements are performed, one straight up, and four with zenith angle  70∘70∘  in directions SW, NW, NE, SE.
    # In ARTS the measurement directions are given by a two-column matrix, where the first column contains the zenith
    # angle and the second column the azimuth angle.
    ws.sensor_los = np.array([[
        0.0,
        0.0,
    ], [70.0, -135.0], [70.0, -45.0], [70.0, 45.0], [70.0, 135.0]])
    ws.sensor_pos = np.array([[2000.0, -21.1, 55.6]] * 5)

    # Reference Measurement
    # Before we can calculate `y`, our setup needs to pass the following tests:
    ws.abs_f_interp_order = 3
    ws.propmat_clearsky_agenda_checkedCalc()
    ws.sensor_checkedCalc()
    ws.atmgeom_checkedCalc()
    ws.atmfields_checkedCalc()
    ws.abs_xsec_agenda_checkedCalc()
    ws.jacobianOff()
    ws.cloudboxOff()
    ws.cloudbox_checkedCalc()

    ws.yCalc()
    y = np.copy(ws.y.value)

    # Setting up the Retrieval
    # In this example, we retrieve ozone and the horizontal and vertical components of the wind velocities.
    # The state space covariance matrix in ARTS is represented by the **covmat_sa** WSV.
    # It belongs to the CovarianceMatrix group, which is used to represent block diagonal matrices.
    # For each retrieval quantity that is added to the retrieval, a corresponding block must be added to **covmat_sa**.
    # This is usually done by the corresponding **retrievalAdd...** call, which looks for this block
    # in the **covmat_block** WSV.
    # In short the general workflow for adding a retrieval quantity is as follows:
    #  - Create the covariance matrix for the retrieval quantity either calling one of the **covmat...** WSV or
    #    by loading your own matrix
    #  - Write the matrix block into **covmat_block**
    #  - Call the **retrievalAdd...** method to add the retrieval quantity and the covariance matrix block
    #    to **covmat_sa**
    lat_ret_grid = np.array([np.mean(ws.lat_grid)])
    lon_ret_grid = np.array([np.mean(ws.lon_grid)])
    n_p = ws.p_grid.value.size

    ws.retrievalDefInit()
    ws.covmat1D(
        ws.covmat_block,
        grid_1=z_grid,
        sigma_1=0.1 * np.ones(n_p),  # Relative uncertainty
        cls_1=10e3 * np.ones(n_p),  # 10km correlation length
        fname="lin")
    ws.retrievalAddAbsSpecies(species="O3",
                              unit="rel",
                              g1=ws.p_grid,
                              g2=lat_ret_grid,
                              g3=lon_ret_grid)
    # Wind u-component
    ws.covmat1D(
        ws.covmat_block,
        grid_1=z_grid[::2],
        sigma_1=100.0 * np.ones(n_p // 2),  # Relative uncertainty
        cls_1=10e3 * np.ones(n_p // 2),  # 10km correlation length
        fname="lin")
    ws.retrievalAddWind(g1=ws.p_grid.value[::2],
                        g2=np.array([np.mean(ws.lat_grid)]),
                        g3=np.array([np.mean(ws.lon_grid)]),
                        component="u")
    # Wind v-component
    ws.covmat1D(
        ws.covmat_block,
        grid_1=z_grid[::2],
        sigma_1=100.0 * np.ones(n_p // 2),  # Relative uncertainty
        cls_1=10e3 * np.ones(n_p // 2),  # 10km correlation length
        fname="lin")
    ws.retrievalAddWind(g1=ws.p_grid.value[::2],
                        g2=np.array([np.mean(ws.lat_grid)]),
                        g3=np.array([np.mean(ws.lon_grid)]),
                        component="v")
    ws.retrievalDefClose()
    ws.covmatDiagonal(ws.covmat_block,
                      ws.covmat_inv_block,
                      vars=0.0001 * np.ones(ws.y.value.shape))
    ws.covmat_seSet(ws.covmat_block)

    @arts_agenda
    def inversion_iterate_agenda(ws):
        ws.x2artsStandard()
        ws.atmfields_checkedCalc()
        ws.atmgeom_checkedCalc()
        ws.yCalc()
        ws.Print(ws.y)
        ws.Print(ws.jacobian)
        ws.VectorAddVector(ws.yf, ws.y, ws.y_baseline)
        ws.IndexAdd(ws.inversion_iteration_counter,
                    ws.inversion_iteration_counter, 1)

    ws.Copy(ws.inversion_iterate_agenda, inversion_iterate_agenda)

    # A Priori State
    # For the a priori state we assume zero wind in any direction. The a priori vector for the OEM is created by
    # the `xaStandard` WSM, which computes $x_a$ from the current atmospheric state.
    ws.wind_u_field.value[:] = 0.0
    ws.wind_v_field.value[:] = 0.0
    ws.xaStandard()

    # The OEM Calculation
    ws.x = np.zeros(0)
    ws.jacobian = np.zeros((0, 0))
    ws.y.value[:] = y
    ws.OEM(method="lm",
           max_iter=20,
           display_progress=1,
           lm_ga_settings=np.array([100.0, 2.0, 2.0, 10.0, 1.0, 1.0]))
    ws.x2artsStandard()

    z = ws.z_field.value[:, 0, 0].ravel()
    wind_u = ws.wind_u_field.value[z > 40e3, 0, 0]
    wind_v = ws.wind_v_field.value[z > 40e3, 0, 0]
    assert np.allclose(wind_u, u_wind, atol=1)
    assert np.allclose(wind_v, v_wind, atol=1)