colindex2 2.10.5__py3-none-any.whl

colindex2/__init__.py

@@ -0,0 +1,8 @@
+from .colindex2 import Detect
+from .commands import detect, track, gen_data_settings, find_track
+from .tracking_overlap2 import Track
+from .get_IDfile import Finder
+from .draw_map import _draw_map, draw_ro
+
+from .modules import gcd
+from .modules import invert_gcd2

colindex2/circle_overlap.py

@@ -0,0 +1,115 @@
+#!/usr/bin/env python3
+import numpy as np
+import pandas as pd
+
+
+def area_heron_x2(a, b, c):
+    s = (a+b+c)/2
+    if min(s-a, s-b, s-c) <= 0:
+        return np.nan
+    else:
+        return np.sqrt(s*(s-a)*(s-b)*(s-c))
+
+
+class circle:
+    def __init__(self, r, x):
+        self.r = r
+        self.x = x
+        self.S = np.pi * r * r
+
+
+class circle_overlap:
+    def __init__(self, c1, c2, d):
+
+        fct = 1
+        S_square = area_heron_x2(c1.r, c2.r, d)
+
+        if np.isnan(S_square):   # cannot make inner triangle
+
+            alpha1 = np.nan
+            d1 = np.nan
+            h = np.nan
+
+            if d < c1.r:
+                overlap = 1.
+            else:
+                overlap = 0.
+
+        else:  # can make inner triangle
+            # https://tjkendev.github.io/procon-library/python/geometry/circles_intersection_area.html
+
+            theta1 = np.arccos((c1.r**2 + d**2 - c2.r**2) / (2 * c1.r * d))
+            theta2 = np.arccos((c2.r**2 + d**2 - c1.r**2) / (2 * c2.r * d))
+
+            S_fan = c1.r**2 * theta1 + c2.r**2 * theta2
+            S_and = S_fan - S_square
+
+            # complete overlap ratio
+            S_or = c1.S + c2.S - S_and
+            overlap = min(S_and / S_or, 1.0)
+
+            # smaller-side overlap ratio
+            #S_small = min(c1.S, c2.S)
+            #overlap = min(S_and / S_small, 1.0)
+
+        self.overlap = overlap
+
+
+def test():
+
+    import matplotlib.pyplot as plt
+    import matplotlib.patches as patches
+
+    r1, r2 = 500, 600
+
+    ds = np.arange(100, 1151, 50)
+    overs = []
+    alpha1s = []
+    d1s = []
+    hs = []
+
+    for d in ds:
+
+        c1 = circle(r=r1, x=0)
+        c2 = circle(r=r2, x=d)
+        o = cirlce_overlap(c1, c2, d)
+        overs.append(o.overlap)
+        alpha1s.append(o.alpha1)
+        d1s.append(o.d1)
+        hs.append(o.h)
+
+
+    for i, d in enumerate(ds):
+
+        print(f'{d=}')
+        
+        ax = plt.axes()
+        C1 = patches.Circle(xy=(0, 500), radius=r1, fc='none', ec='tab:blue')
+        C2 = patches.Circle(xy=(d, 500), radius=r2, fc='none', ec='tab:blue')
+        ax.add_patch(C1)
+        ax.add_patch(C2)
+        ax.plot(0, 500, marker='.', c='tab:blue')
+        ax.plot(d, 500, marker='.', c='tab:blue')
+        ax.plot([d1s[i], d1s[i]], [500-hs[i], 500+hs[i]], c='r')
+        ax.set_aspect('equal')
+        ax.set_xlim([-700,1250])
+        ax.set_ylim([-200, 1200])
+        ax.set_title(f'O={overs[i]:.3f}')
+        plt.savefig(f'fig/cc_d{d}.png')
+        plt.close()
+
+        ax = plt.axes()
+        ax.plot(ds, overs)
+        ax.plot(ds[i], overs[i], '.r')
+        plt.savefig(f'fig/overs_{d}.png')
+        plt.close()
+
+        '''ax = plt.axes()
+        ax.plot(ds, np.array(alpha1s)/np.pi)
+        ax.plot(ds[i], alpha1s[i]/np.pi, '.r')
+        plt.savefig(f'fig/alpha1s_d{d}.png')
+        plt.close()'''
+        
+
+if __name__ == '__main__':
+    test()

colindex2/colindex2.py

@@ -0,0 +1,319 @@
+import os
+import numpy as np
+import pandas as pd
+import xarray as xr
+#import multiprocessing as mp
+import warnings 
+warnings.filterwarnings('ignore') 
+
+from .modules import check_gpu
+from .modules import _calc_as_torch
+
+class Detect():
+    """
+    Depression detector by searching each optimal location, intensity, and size from 2D fields.
+    See puplication for details (Kasuga et al. 2021; K21) and (Kasuga and Honda 2025; KH25).
+
+    ----
+
+    Parameters:
+    -----------
+
+    da : xarray.DataArray or numpy.ndarray
+        The target 2D field with (or without) time axis.  
+        The dimmension order must be time->latitude->longitude.
+
+    odir : str, default './d01'
+        Parent output directory path. By defaults, output data will be stored as follows
+        Outputs will be produced in {odir}/AS (AS field) and {odir}/V (detected point values).
+
+    ty : str, {'both', 'L', 'H'}, default 'both'
+        Type of vortices to be detected.
+
+    r : np.array, default np.arange(200, 2101, 100) (K21)
+        Radial searching variable for average slop function (AS) [km].
+        Set range to include scales of which phenomena you are interested in.
+
+    SR_thres : float, default 3. (K21)
+        Threshold to remove noises with respect to the slope ratio.
+
+    So_thres : float, default 3. (KH25)
+        Threshold to remove noises with respect to the intensity [m/100km].
+        In K21, ``10`` and ``5`` was used for climatology for cutoff lows and preexisting troughs, respectively.
+
+    Do_thres : float, default 0.
+        Threshold to remove noises with respect to the depth [m].
+        A large value (e.g., 10.) reduces small-size noise, which might be helpful for surface low/high detection.
+
+    xx_thres : float, default 0.5
+        Threshold to remove noises with respect to the zonal laplacan [m/(100 km)^2].
+        This works for high detection
+        by removeing large and weak ridges located south of the subtropical jet.
+
+    rm_rmin : bool, default True
+        Detected point will be removed if ro is minimum of r because this may be too small.
+        True in K21 and KH25.
+
+    rm_rmax : bool, default False
+        Detected point will be removed if ro is maximum of r because this may be too large.
+        True in K21 and False in KH25.
+
+    factor : float, default 100.
+        Factor to convert units of So and AS. Default convert m/km -> m/100 km for geopotential height input.
+
+    local_ex_req_num : int, default 7.
+        Number required in local extremum condition (N). 
+        The conditions of minimum/maximum are defined in a 3x3 grid box as at least N points (out of 8 points)
+        must greater/smaller than and equal to the center value.  
+
+    gpu : bool, default False
+        Note, this feature is experimental.
+        If True, the processes wil be computed with gpu processing via CUDA or METAL.
+
+    nc : bool, default True
+        If True, netcdf file will be created for AS. If you want grads binary outputs, set False.
+
+    fmt : str, default '<f4'
+        Output binary format when nc=False. The format is correspond to dtype in numpy.
+
+    subs : bool, default False
+        If True, 2D fields of ro, m, and n are produced.
+
+    t : pandas.DatetimeIndex, default [pd.Timestamp('1000-01-01 00:00')]
+        Time of data. Option for numpy.ndarray input.
+
+    lev : int, default 0
+        Level of data. Option for numpy.ndarray input. This will be used in the name of output files.
+        Note: Float level will be converted to integer.
+
+    lons : list or numpy.ndarray, default [-999.]
+        Longitude list. Option for numpy.ndarray input. Set like lon=np.arange(0, 360, 1.25).
+
+
+    lats : list or numpy.ndarray, default [-999.]
+        Option for numpy.ndarray input. Set like lon=np.arange(-90, 90, 1.25).
+
+    """
+
+    def __init__(self,
+                 da,  # 2D DataArray or ndarray
+                 odir='./d01',  # output parent dir
+                 ty="both",
+                 r=np.arange(200, 2101, 100),  # searching radius variable
+                 SR_thres=3.,  # shape threshold
+                 So_thres=3.,  # intensity threshold
+                 Do_thres=0.,  # depth threshold
+                 xx_thres=.5,  # zonal concavity thershold
+                 rm_rmin=True,  # remove minimum size
+                 rm_rmax=False,  # remove maximum size
+                 factor=100.,  # factor of So m/km -> m/100km
+                 local_ex_req_num=7, # out of 8 points
+                 gpu=False, # gpu calculation
+                 nc=True,  # netcdf outputs otherwise grads bin
+                 fmt='<f4',  # binary format when nc=False
+                 subs=False, # output intermediate fields
+                 lev=0,  # for ndarray input below
+                 t=[pd.Timestamp('1700-01-01 00:00')],
+                 lons=[-999.],
+                 lats=[-999.],
+                 ):
+
+        # constants
+        self.rr = 6371.  # earth radius [km]
+        self.g = 9.8  # gravity [m/ss]
+        self.factor = factor
+
+        self.ty = ty
+        self.r = np.array(r)
+
+        self.da = da
+        self.odir = odir
+        self.SR_thres = SR_thres
+        self.So_thres = So_thres
+        self.Do_thres = Do_thres
+        self.xx_thres = xx_thres
+        self.thres4 = [So_thres, Do_thres, SR_thres, xx_thres]
+        self.rm_rmin = rm_rmin
+        self.rm_rmax = rm_rmax
+        self.local_ex_req_num = local_ex_req_num
+
+        self.lons = lons
+        self.lats = lats
+        self.t = t
+        self.lev = lev
+
+        self.gpu = gpu
+        self.nc = nc
+        self.fmt = fmt
+        self.subs = subs
+
+        # main process
+        self._main()
+
+        # maxe a text about detection parameters
+        self._save_param_text()
+
+        print("done!")
+
+
+    def _set_coords(self):
+
+        da = self.da
+        lons = self.lons
+        lats = self.lats
+        lev = self.lev
+        t = self.t
+        if isinstance(da, xr.DataArray):
+
+            _dims = [k for k, v in da.coords.items()]  # this may include squeezed dims
+
+            if 'valid_time' in _dims:
+                tt = pd.to_datetime(da.valid_time.values)
+            elif 'time' in _dims:
+                tt = pd.to_datetime(da.time.values)
+
+            if isinstance(tt, pd.Timestamp):
+                self.t = [tt]  # must be iterable
+                self.da = da.values[np.newaxis, ...]
+            else:
+                self.t = tt  # pd.DatetimeIndex or default list
+                self.da = da.values
+
+            if 'level' in _dims:
+                self.lev = da.level.values
+            elif 'lev' in _dims:
+                self.lev = da.lev.values
+            elif 'levs' in _dims:
+                self.lev = da.levs.values
+            elif 'isobaricInhPa' in _dims:
+                self.lev = da.isobaricInhPa.values
+
+            if 'longitude' in _dims:
+                self.lon = da.longitude.values
+            if 'lon' in _dims:
+                self.lon = da.lon.values
+            if 'lons' in _dims:
+                self.lon = da.lons.values
+
+            if 'latitude' in _dims:
+                self.lat = da.latitude.values
+            if 'lat' in _dims:
+                self.lat = da.lat.values
+            if 'lats' in _dims:
+                self.lat = da.lats.values
+
+        elif isinstance(da, np.ndarray):
+
+            if len(lons) != da.shape[-1] or len(lats) != da.shape[-2]:
+                raise ValueError('shape of lons or lats not fit input')
+            else:
+                self.t = t
+                self.lev = lev
+                self.lon = lons
+                self.lat = lats
+                self.da = da
+
+            if len(self.t) == 1:
+                self.da = [self.da]
+                self.t = [self.t]
+
+        else:
+            raise ValueError('wrong input data type, must be xr.DataArray or np.ndarray')
+
+        if isinstance(self.lev, np.ndarray):
+            self.lev = self.lev.item()
+
+        self.lev = int(self.lev)
+
+
+
+    def check_zonal_loop(self):
+
+        lon1 = self.lon[1]
+        lon1n = self.lon[-1]
+        lon0 = self.lon[0]
+        dlon = lon1 - lon0
+
+        if lon1n == lon0:
+            raise ValueError("longitude is already looped. use non-looped data")
+
+        l = lon1n + dlon
+
+        if l == lon0 or l+360 == lon0 or l-360 == lon0: 
+            self.zloop = True
+        else:
+            self.zloop = False
+
+    def check_increasing_lat(self):
+        
+        dlat = self.lat[1] - self.lat[0]
+
+        if dlat < 0.:
+            self.da = self.da[...,::-1,:].copy()
+            self.lat = self.lat[::-1].copy()
+
+
+    def _main(self):
+
+        # set coordinates of data
+        self._set_coords()
+
+        # check gpu is available
+        self.device = check_gpu(self.gpu)
+
+        # check whether data longitude is looped
+        # and increasing latitude
+        self.check_zonal_loop()
+        self.check_increasing_lat()
+
+        print("level:", self.lev)
+        print("calculation:", self.device)
+        print("-----------------------")
+
+        # start main process
+        args = [(
+            self.odir,
+            self.ty,
+            self.da[i],
+            self.lon,
+            self.lat,
+            self.lev,
+            self.t[i],
+            self.r,
+            self.thres4,
+            self.rm_rmin,
+            self.rm_rmax,
+            self.factor,
+            self.zloop,
+            self.device,
+            self.nc,
+            self.fmt,
+            self.subs,
+            self.local_ex_req_num,
+            ) for i in range(len(self.t))]
+
+        #with mp.Pool(self.nproc) as p:
+        #    p.map(_calc_as_torch, args)
+        for x in args:
+            _calc_as_torch(x)
+
+    def _save_param_text(self):
+        from datetime import datetime
+        import getpass
+        import socket
+        now = datetime.now()
+        user = getpass.getuser()
+        host = socket.gethostname()
+        with open(f"{self.odir}/_detection_parms.txt", "w", encoding="utf-8") as f:
+            print("r:", self.r, file=f)
+            print("SR_thres:", self.SR_thres, file=f)
+            print("So_thres:", self.So_thres, file=f)
+            print("Do_thres:", self.Do_thres, file=f)
+            print("xx_thres:", self.xx_thres, file=f)
+            print("rm_rmin:", self.rm_rmin, file=f)
+            print("rm_rmax:", self.rm_rmax, file=f)
+            print("factor:", self.factor, file=f)
+            print("local_ex_req_num:", self.local_ex_req_num, file=f)
+            print("calculation device:", self.device, file=f)
+            print("", file=f)
+            print("executed in", now, "by", user, "@", host, file=f)

colindex2/commands.py

@@ -0,0 +1,269 @@
+import sys, os
+import argparse
+import importlib.util
+import numpy as np
+import pandas as pd
+import xarray as xr
+from .colindex2 import Detect
+from .tracking_overlap2 import Track
+from .get_IDfile import Finder
+import shutil
+from pathlib import Path
+import importlib.resources
+
+def gen_data_settings():
+    with importlib.resources.path("colindex2", "data_settings.py") as file_path:
+        dest = Path.cwd() / "data_settings.py"
+        shutil.copy(file_path, dest)
+
+def _load_local_module(path):
+    if not os.path.exists(path):
+        raise FileNotFoundError(f"{path} has not generated yet. Use gen_data_settings command")
+    spec = importlib.util.spec_from_file_location("data_settings", path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules["data_settings"] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def detect():
+    description = """
+Run detection of colindex.
+Before run, generate data_settings.py with 'gen_data_settings' command and edit it.
+
+ $ gen_data_settings
+
+Then execute below for netcdf input.
+
+ $ detect path/to/file.nc
+
+For binary input, specify start/end times (st,et) and frequency of timesteps (freq).
+
+ $ detect path/to/file.grd -st "yyyy-mm-dd hh" -et "yyyy-mm-dd hh" -freq 6h
+
+Option -track execute also tracking after detection.
+In this case, 'tracking_times' in data_settings.py is ignored,
+and the detection timesteps will be used.
+
+ $ detect path/to/file.nc -track
+ $ detect path/to/file.grd -st "yyyy-mm-dd hh" -et "yyyy-mm-dd hh" -freq 6h -track
+
+Result data will be saved in output_dir/V output_dir/AS.
+            """
+#Option -n can set number of multiprocess (default 4)
+#
+# $ detect path/to/file.nc -n 20
+
+    parser = argparse.ArgumentParser(description=description,
+                                     formatter_class=argparse.RawTextHelpFormatter)
+    parser.add_argument("input_data_path",
+                        type=str,
+                        help="input data path")
+    #parser.add_argument("-n",
+    #                    type=int,
+    #                    help="number of multiprocessing, default is 4.",
+    #                    default=4)
+    parser.add_argument("-st",
+                        type=str,
+                        help="start timestep yyyy-mm-dd hh",
+                        default="1700-01-01 00")
+    parser.add_argument("-et",
+                        type=str,
+                        help="end timestep yyyy-mm-dd hh",
+                        default="1700-01-01 00")
+    parser.add_argument("-freq",
+                        type=str,
+                        help="frequency of timestep for pd.date_range (default: '6h')",
+                        default="6h")
+    parser.add_argument("-track",
+                        action="store_true",
+                        help="enable tracking",)
+    args = parser.parse_args()
+    input_data_path = args.input_data_path
+    st = args.st
+    et = args.et
+    freq = args.freq
+    _track = args.track
+
+    data_settings_path = os.path.join(os.getcwd(), "data_settings.py")
+    if not os.path.exists(data_settings_path):
+        raise FileNotFoundError("No data_settings.py in the current dir. Run 'gen_data_settings' command.")
+    D = _load_local_module(data_settings_path)
+
+    if D.input_data_type == "nc":
+        if not hasattr(D, "var_name"):
+            da = xr.open_dataarray(input_data_path)
+        else:
+            ds = xr.open_dataset(input_data_path)
+            da = ds[D.var_name]
+
+        if hasattr(D, "lon_name"):
+            da = da.rename({D.lon_name:"longitude"})
+        if hasattr(D, "lat_name"):
+            da = da.rename({D.lat_name:"latitude"})
+        if hasattr(D, "lev_name"):
+            da = da.rename({D.lev_name:"level"})
+        if hasattr(D, "time_name"):
+            da = da.rename({D.time_name:"time"})
+
+        lons = da.longitude.values
+        lats = da.latitude.values
+        levs = da.level.values
+        times = pd.to_datetime(da.time.values)
+        da = da.values
+        fmt = None
+
+    else:
+        lons = D.lons
+        lats = D.lats
+        levs = D.levs
+        #times = D.times
+        times = pd.date_range(st,et,freq=freq)
+        shape = (len(times),len(levs),len(lats),len(lons))
+        fmt = D.input_data_type
+        with open(input_data_path, "r") as a:
+            da = np.fromfile(a, dtype=fmt).reshape(shape)
+
+    if D.output_data_type == "nc":
+        nc = True
+    else:
+        nc = False
+        fmt = D.output_data_type
+
+    idx_levs = np.where(np.isin(levs, D.selected_levs))[0]
+
+    for l in idx_levs:
+        Detect(da[:,l,:,:],
+               D.output_dir,
+               D.detection_type,
+               D.r,
+               D.SR_thres,
+               D.So_thres,
+               D.Do_thres,
+               D.xx_thres,
+               D.rm_rmin,
+               D.rm_rmax,
+               D.factor,
+               D.local_ex_req_num,
+               D.gpu,
+               nc,
+               fmt,
+               D.subs,
+               levs[l],
+               times,
+               lons,
+               lats,
+               )
+
+    if _track:
+        track_cli(times)
+
+
+def track_cli(times=0):
+
+    data_settings_path = os.path.join(os.getcwd(), "data_settings.py")
+    if not os.path.exists(data_settings_path):
+        raise FileNotFoundError("No data_settings.pn in the current dir. Run 'gen_data_settings' command.")
+    D = _load_local_module(data_settings_path)
+
+    if isinstance(times, pd.DatetimeIndex):
+        tracking_times = times
+    else:
+        tracking_times = D.tracking_times
+
+    argss=[(
+        D.output_dir,
+        D.tracking_types,
+        lev,
+        tracking_times,
+        D.tlimit,
+        D.penalty,
+        D.long_term,
+        D.operational,
+        D.DU_thres,
+        D.QS_min_intensity,
+        D.QS_min_radius,
+        D.QS_min_overlap_ratio,
+        D.only_count
+        ) for lev in D.selected_levs]
+
+    for _args in argss:
+        Track(*_args)
+
+
+def track():
+    description = """
+Run tracking of colindex. Before run, edit TRACKING SETTINGS section in data_settings.py
+
+ $ track
+
+Result data will be saved in output_dir/Vtc.
+"""
+    parser = argparse.ArgumentParser(description=description,
+                                     formatter_class=argparse.RawTextHelpFormatter)
+    parser.parse_args()
+    track_cli()
+
+
+def find_track():
+    description = """
+Search tracking for specific IDs and save. Four arguments are required.
+ $ find_track output_dir type lev freqH type2
+
+The fourth argument 'type2' can be set 3 types:
+
+c: only count number of all tracks
+ $ find_track ./d01 L 300 6 c
+
+A: find all tracks and save them as an all-included csv output_dir/all_{type}_{lev}.csv.
+ $ find_track ./d01 L 300 6 A
+
+a: find all tracks and save them separately in output_dir/ID/.
+ $ find_track ./d01 L 300 6 a
+
+(digit of ID): find one track labeled by ID. Add fifth argument of 'yyyymm' (year month) when the track observed, as following
+ $ find_track ./d01 L 300 6 2222 202302
+
+ This will produce a track obtained from 6-hour intarval data for 300-hPa level whose ID is 2222, appears in Feb 2023. 
+ If add -a (-b) option at its tail, the tracks after merging (before splitting) are connected.
+ $ find_track ./d01 L 300 6 2222 202302 -a
+ $ find_track ./d01 L 300 6 2222 202302 -b
+ $ find_track ./d01 L 300 6 2222 202302 -a -b
+
+ """
+    parser = argparse.ArgumentParser(description=description, formatter_class=argparse.RawTextHelpFormatter)
+    parser.add_argument("odir", type=str, help="output data directory")
+    parser.add_argument("ty", type=str, help="L or H")
+    parser.add_argument("lev", type=str, help="level")
+    parser.add_argument("freqH", type=int, help="frequency of timestep in hour")
+    parser.add_argument("type2", nargs="+", type=str, help="A or a or c or 'ID yyyymm'")
+    parser.add_argument('-a', '--option_a', action='store_true', default=False, help='follow track after merge')
+    parser.add_argument('-b', '--option_b', action='store_true', default=False, help='follow track before split')
+    args = parser.parse_args()
+    odir = args.odir
+    ty = args.ty
+    #lev = int(args.lev)
+    lev = args.lev
+    type2 = args.type2
+
+    #data_settings_path = os.path.join(os.getcwd(), "data_settings.py")
+    #D = _load_local_module(data_settings_path)
+    #freqH = int(D.tracking_times.inferred_freq[:-1])
+    freqH = args.freqH
+
+    f = Finder(odir, timestep=freqH)
+
+    if type2[0] == "a":
+        f.find_all(ty,lev)
+    elif type2[0] == "c":
+        f.count_all(ty,lev)
+    elif type2[0] == "A":
+        f.find_all(ty,lev,all_in_one=True)
+    elif args.option_a and args.option_b:
+        f.find_one(ty,lev,int(type2[1]), int(type2[0]),after_merge=True,before_split=True)
+    elif args.option_a:
+        f.find_one(ty,lev,int(type2[1]), int(type2[0]),after_merge=True)
+    elif args.option_b:
+        f.find_one(ty,lev,int(type2[1]), int(type2[0]),before_split=True)
+    else:
+        f.find_one(ty,lev,int(type2[1]), int(type2[0]))