-
Notifications
You must be signed in to change notification settings - Fork 1
/
create_test_data.py
139 lines (117 loc) · 4.3 KB
/
create_test_data.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
"""
Creates data objects for testing
"""
import sys
import os
import numpy as np
import xarray as xr
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(SCRIPT_DIR))
from utils.create_artificial_inventories import ArtificialInventory
def create_test_conc_resp():
"""
Creates an example response dataset for testing purposes
with resp_type = "conc"
Returns:
xr.Dataset: A minimal response dataset with random data.
"""
lat_arr = np.arange(-75.0, 80.0, 15.0).astype("float32")
dim_lat = len(lat_arr)
plev_arr = np.arange(100.0, 1000.0, 100.0).astype("float32")
dim_plev = len(plev_arr)
emi_lat_arr = np.array([10.0, 40.0], dtype="float32")
emi_plev_arr = np.array([250.0, 500.0], dtype="float32")
emi_loc_arr = np.array([["p10_250", "p10_500"], ["p40_250", "p40_500"]])
p10_250_arr = np.random.randn(dim_lat, dim_plev).astype("float32")
p40_250_arr = np.random.randn(dim_lat, dim_plev).astype("float32")
p10_500_arr = np.random.randn(dim_lat, dim_plev).astype("float32")
p40_500_arr = np.random.randn(dim_lat, dim_plev).astype("float32")
resp = xr.Dataset(
data_vars={
"emi_loc": (["emi_lat", "emi_plev"], emi_loc_arr),
"p10_250": (["lat", "plev"], p10_250_arr),
"p40_250": (["lat", "plev"], p40_250_arr),
"p10_500": (["lat", "plev"], p10_500_arr),
"p40_500": (["lat", "plev"], p40_500_arr),
},
coords={
"lat": lat_arr,
"plev": plev_arr,
"emi_lat": emi_lat_arr,
"emi_plev": emi_plev_arr,
},
attrs={"resp_type": "conc"},
)
return resp
def create_test_rf_resp():
"""
Creates an example response dataset for testing purposes
with resp_type = "rf"
Returns:
xr.Dataset: A minimal response dataset with random data.
"""
emi_lat_arr = np.array([10.0, 40.0], dtype="float32")
emi_plev_arr = np.array([250.0, 500.0], dtype="float32")
emi_loc_arr = np.array([["p10_250", "p10_500"], ["p40_250", "p40_500"]])
h2o_arr = np.random.rand(len(emi_lat_arr), len(emi_plev_arr)).astype(
"float32"
)
emi_air_mass_arr = np.ones_like(h2o_arr)
resp = xr.Dataset(
data_vars={
"emi_air_mass": (["emi_lat", "emi_plev"], emi_air_mass_arr),
"emi_loc": (["emi_lat", "emi_plev"], emi_loc_arr),
"H2O": (["emi_lat", "emi_plev"], h2o_arr),
},
coords={
"emi_lat": emi_lat_arr,
"emi_plev": emi_plev_arr,
},
attrs={"resp_type": "rf"},
)
return resp
def create_test_inv(year=2020, size=3):
"""
Creates an example inventory dataset for testing purposes.
Args:
year (int): inventory year
size (int): The number of samples to generate.
Returns:
xr.Dataset: An xarray dataset with random inventory data.
"""
inv = ArtificialInventory(year, size=size).create()
return inv
def create_test_resp_cont(n_lat=48, n_lon=96, n_plev=39, seed=None):
"""Creates example precalculated contrail input data for testing purposes.
Args:
n_lat (int, optional): Number of latitude values. Defaults to 48.
n_lon (int, optional): Number of longitude values. Defaults to 96.
n_plev (int, optional): Number of pressure level values. Defaults to 39.
seed (int, optional): Random seed.
Returns:
xr.Dataset: Example precalculated contrail input data.
"""
# set random seed
np.random.seed(seed)
# Create the coordinates
lon = np.linspace(0, 360, n_lon, endpoint=False)
lat = np.linspace(90, -90, n_lat + 2)[1:-1] # do not include 90 or -90
plev = np.linspace(1014, 10, n_plev)
# Create the data variables with random values between 0 and 1
iss = np.random.rand(n_lat, n_lon)
sac_con = np.random.rand(n_lat, n_lon, n_plev)
sac_lh2 = np.random.rand(n_lat, n_lon, n_plev)
# Combine into an xarray Dataset
ds_cont = xr.Dataset(
{
"ISS": (["lat", "lon"], iss),
"SAC_CON": (["lat", "lon", "plev"], sac_con),
"SAC_LH2": (["lat", "lon", "plev"], sac_lh2)
},
coords={
"lon": ("lon", lon),
"lat": ("lat", lat),
"plev": ("plev", plev)
}
)
return ds_cont