get_domain.py 2.8 KB
Newer Older
Isabelle Pison's avatar
Isabelle Pison committed
1
2
3
4
5
6
import numpy as np
import glob
import datetime
import os
import xarray as xr
from pycif.utils.classes.setup import Setup
Isabelle Pison's avatar
Isabelle Pison committed
7
import copy
Isabelle Pison's avatar
Isabelle Pison committed
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

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
36

Isabelle Pison's avatar
Isabelle Pison committed
37
38
39
40
41
    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
42
43
44
45
    # must be increasing
    if lat[1] < lat[0]: 
      lat2 = np.flip(lat)
      lat = copy.copy(lat2)
Isabelle Pison's avatar
Isabelle Pison committed
46
    nlon = lon.size
Isabelle Pison's avatar
Isabelle Pison committed
47
48
    nlat = lat.size 
    
Isabelle Pison's avatar
Isabelle Pison committed
49
50
51
52
53
54
    #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
55
    #print(lon_min)
Isabelle Pison's avatar
Isabelle Pison committed
56
    
Isabelle Pison's avatar
Isabelle Pison committed
57
58
59
60
61
62
    # 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
63
    # Read vertical information in domain_file
64
65
66
    sigma_a = nc["hyam"].values
    sigma_b = nc["hybm"].values
    nlevs = sigma_a.size
Isabelle Pison's avatar
Isabelle Pison committed
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
    
    # 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,
                "sigma_a": sigma_a[1:],
                "sigma_b": sigma_b[1:],
                "pressure_unit": "Pa"
            }
        }
    )

    Setup.load_setup(setup, level=1)
Isabelle Pison's avatar
Isabelle Pison committed
92
93
94
95
96
97
98
    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
99
100
101
102
    
    # Cyclic domain in longitude
    setup.domain.lon_cyclic = True
    
Isabelle Pison's avatar
Isabelle Pison committed
103
    return setup.domain