get_domain.py 2.91 KB
Newer Older
Isabelle Pison's avatar
Isabelle Pison committed
1
2
3
4
5
import numpy as np
import glob
import datetime
import os
import xarray as xr
6
from .....utils.classes.setup import Setup
Isabelle Pison's avatar
Isabelle Pison committed
7
import copy
Isabelle Pison's avatar
Isabelle Pison committed
8

9

Isabelle Pison's avatar
Isabelle Pison committed
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
def get_domain(ref_dir, ref_file, input_dates, target_dir, tracer=None):
    
    # Looking for a reference file to read lon/lat in
    list_file = glob.glob("{}/*nc".format(ref_dir))

    domain_file = None
    # Either a file is specified in the Yaml
    if ref_file in list_file:
        domain_file = "{}/{}".format(ref_dir, ref_file)

    # Or loop over available file regarding file pattern
    else:
        for bc_file in list_file:
            try:
                date = datetime.datetime.strptime(
                    os.path.basename(bc_file), ref_file
                )
                domain_file = bc_file
                break
            except ValueError:
                continue

    if domain_file is None:
        raise Exception(
            "CAMS domain could not be initialized as no file was found"
        )
    print('Domain file for CAMS BCs:',domain_file)
Isabelle Pison's avatar
Isabelle Pison committed
37

Isabelle Pison's avatar
Isabelle Pison committed
38
39
40
41
42
    nc = xr.open_dataset(domain_file, decode_times=False)

    # Read lon/lat in domain_file   
    lon = nc["longitude"]
    lat = nc["latitude"]
Isabelle Pison's avatar
Isabelle Pison committed
43
44
45
46
    # must be increasing
    if lat[1] < lat[0]: 
      lat2 = np.flip(lat)
      lat = copy.copy(lat2)
Isabelle Pison's avatar
Isabelle Pison committed
47
    nlon = lon.size
Isabelle Pison's avatar
Isabelle Pison committed
48
49
    nlat = lat.size 
    
Isabelle Pison's avatar
Isabelle Pison committed
50
51
52
53
54
55
    #print(nlon,nlat)
    #print(lon[0])
    lon_min = lon.min()
    lon_max = lon.max() 
    lat_min = lat.min() 
    lat_max = lat.max() 
Isabelle Pison's avatar
Isabelle Pison committed
56
    #print(lon_min)
Isabelle Pison's avatar
Isabelle Pison committed
57
    
Isabelle Pison's avatar
Isabelle Pison committed
58
59
60
61
62
63
    # Compute corners #WARNING: only valid for regular grid in lat/lon -> to generalize
    dx = (lon[1] - lon[0]) / 2
    dy = (lat[1] - lat[0]) / 2
    lonc = np.linspace(lon_min - dx/2., lon_max + dx/2., nlon + 1)
    latc = np.linspace(lat_min - dy/2., lat_max + dy/2., nlat + 1)

Isabelle Pison's avatar
Isabelle Pison committed
64
    # Read vertical information in domain_file
Elise Potier's avatar
Elise Potier committed
65
66
67
68
69
70
    if tracer.aibi_name:
         sigma_a = nc["ap"].values
         sigma_b = nc["bp"].values
    else :
         sigma_a = nc["hyam"].values
         sigma_b = nc["hybm"].values
71
        
72
    nlevs = sigma_a.size
73

Isabelle Pison's avatar
Isabelle Pison committed
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
    # Initializes domain
    setup = Setup.from_dict(
        {
            "domain": {
                "plugin": {
                    "name": "dummy",
                    "version": "std",
                    "type": "domain",
                },
                "xmin": lon_min,
                "xmax": lon_max,
                "ymin": lat_min,
                "ymax": lat_max,
                "nlon": nlon,
                "nlat": nlat,
                "nlev": nlevs,
90
91
                "sigma_a": sigma_a,
                "sigma_b": sigma_b,
Isabelle Pison's avatar
Isabelle Pison committed
92
93
94
95
96
97
                "pressure_unit": "Pa"
            }
        }
    )

    Setup.load_setup(setup, level=1)
Isabelle Pison's avatar
Isabelle Pison committed
98
99
100
101
102
103
104
    zlon, zlat = np.meshgrid(lon, lat)
    zlonc, zlatc = np.meshgrid(lonc, latc)

    setup.domain.zlon = zlon
    setup.domain.zlat = zlat
    setup.domain.zlonc = zlonc
    setup.domain.zlatc = zlatc
105
106
107
108
    
    # Cyclic domain in longitude
    setup.domain.lon_cyclic = True
    
Isabelle Pison's avatar
Isabelle Pison committed
109
    return setup.domain