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]
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