Heatwaves/intern/support.py at main · thigs0/Heatwaves · GitHub

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
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
from alive_progress import alive_bar
import xarray as xr
import numpy as np
import pandas as pd
import warnings

def heatwave(dataset_tmax:xr.Dataset, opt:int, n:int, dataset_tmin=xr.Dataset(), percent_tmax=xr.Dataset(), percent_tmin=xr.Dataset(),
             yearmin=1900, with_anomaly=False, with_season=False) -> xr.Dataset:
    match opt:
        case 1:
            #reduce the size to only the interval analized
            dataset_tmax = dataset_tmax.sel(time = dataset_tmax.time[dataset_tmax.time.dt.year > yearmin])
            greater = hotdays_opt1(dataset_tmax, percent_tmax)
            if with_anomaly:
                return xr.Dataset()
            else:
                return heatwave_opt1(greater, n)
        case 2:
            #reduce the size to only the interval analized
            dataset_tmax = dataset_tmax.sel(time = dataset_tmax.time[dataset_tmax.time.dt.year > yearmin])
            dataset_tmin = dataset_tmin.sel(time = dataset_tmin.time[dataset_tmin.time.dt.year > yearmin])
            greater = hotdays_opt2(dataset_tmax, percent_tmax, dataset_tmin, percent_tmin)
            if with_anomaly:
                return heatwave_opt2_with_anomaly(greater, dataset_tmax, percent_tmax, n)
            else:
                return heatwave_opt2(greater, n)
        case 3:
            return xr.Dataset()
        case 4:
            greater = hotdays_opt1(dataset_tmax, percent_tmax)
            if with_anomaly:
                return xr.Dataset()
            else:
                return heatwave_opt4(greater)
        case _:
            raise ValueError("Invalid value")

def hotdays_opt2(ds, ds_norm, ds2, ds2_norm) -> xr.Dataset:
    out = ds.copy()
    out['tmin'] = ds2.tmin
    # Initialize heatwave variable with False values
    out['greater'] = xr.zeros_like(ds.tmax, dtype=bool)

    # Percorre cada data do dataset original
    with alive_bar(len(ds.time)) as bar:
        for i, date in enumerate(ds.time):
            # export the date value
            this_date = date.values
            month = date.dt.month.item()
            day = date.dt.day.item()

            # current date
            ds1_current = ds.sel(time=this_date)
            ds2_current = ds2.sel(time=this_date)

            # Valor normal para aquele dia/mês
            ds1_nor = ds_norm.sel(time=((ds_norm.time.dt.month == month) &
                                        (ds_norm.time.dt.day == day))).isel(time=0)
            ds2_nor = ds2_norm.sel(time=((ds_norm.time.dt.month == month) &
                                        (ds_norm.time.dt.day == day))).isel(time=0)

            # Calculate heatwave condition - FIXED LINE
            heatwave_condition = ((ds1_current.tmax > ds1_nor.tmax) &
                                 (ds2_current.tmin > ds2_nor.tmin))
            # Assign only to the heatwave variable - FIXED LINE
            out['greater'].loc[dict(time=this_date)] = heatwave_condition
            bar() # Update the progress bar
    return out

def hotdays_opt1(ds, ds_norm):
    out = ds.copy()
    out['greater'] = xr.zeros_like(ds.tmax, dtype=bool)

    # Percorre cada data do dataset original
    with alive_bar(len(ds.time)) as bar:
        for i, date in enumerate(ds.time):
            # Extrai o valor da data
            this_date = date.values
            month = date.dt.month.item()
            day = date.dt.day.item()

            # Valor do dia atual
            current_val = ds.sel(time=this_date)

            # Valor normal para aquele dia/mês
            normal_val = ds_norm.sel(time=((ds_norm.time.dt.month == month) &
                                        (ds_norm.time.dt.day == day))).isel(time=0)

            out['greater'].loc[dict(time=this_date)] = (current_val["tmax"] > normal_val["tmax"]).astype(int)
            bar() # Update the progress bar
    return out

def heatwave_opt2(ds:xr.Dataset, n:int):
    # Cria uma cópia booleana com False por padrão
    ds = ds.astype(int)
    out = ds.copy()

    # Initialize temporal variables as boolean arrays
    out['heatwave'] = xr.zeros_like(ds.greater, dtype=bool)

    # Percorre cada data do dataset original
    len_date = len(ds.time[:-n+1])-1
    with alive_bar(len_date * len(ds.lat.values) * len(ds.lon.values)) as bar:
        for lat in ds.lat.values:
            for lon in ds.lon.values:
                heat = np.zeros(len_date+n)
                i = 0
                while i < len_date:
                    value = ds.greater.sel(time=ds.time.values[i:i + n + 1], lat=lat, lon=lon)
                    if sum(value.values) == n:
                        heat[i] = 1
                        i+=3
                        bar(3) # Update the progress bar
                    else:
                        i+=1
                        bar() # Update the progress bar
                out.heatwave.loc[dict(lat=lat, lon=lon)] = heat
    return out

def heatwave_opt4(ds:xr.Dataset):
    # Cria uma cópia booleana com False por padrão
    ds = ds.astype(int)
    out = ds.copy()

    # Initialize temporal variables as boolean arrays
    out['heatwave'] = xr.zeros_like(ds.greater, dtype=bool)
    out['extention'] = xr.zeros_like(ds.greater, dtype=bool)

    # Percorre cada data do dataset original
    len_date = len(ds.time[:-3+1])-1
    with alive_bar(len_date * len(ds.lat.values) * len(ds.lon.values)) as bar:
        for lat in ds.lat.values:
            for lon in ds.lon.values:
                heat = np.zeros(len_date+3)
                extention = np.zeros(len_date+3)
                i = 0
                while i < len_date:
                    value = ds.greater.sel(time=ds.time.values[i:i + 3 + 1], lat=lat, lon=lon)
                    if sum(value.values) == 3: # begin a heatwave
                        k=i+3+1
                        while k < len_date and ds.greater.sel(time=ds.time.values[k], lat=lat, lon=lon) == 1:
                            k+=1
                        heat[i] = 1
                        extention[i] = k
                        i=k-1
                        bar(k-1) # Update the progress bar
                    else:
                        i+=1
                        bar() # Update the progress bar
                out.heatwave.loc[dict(lat=lat, lon=lon)] = heat
                out.extention.loc[dict(lat=lat, lon=lon)] = extention
    return out

def heatwave_opt1(ds:xr.Dataset, n:int):
    ds = ds.astype(int)
    out = ds.copy()

    # Initialize temporal variables as boolean arrays
    out['heatwave'] = xr.zeros_like(ds.greater, dtype=bool)

    len_date = len(ds.time[:-n+1])-1
    with alive_bar(len_date * len(ds.lat.values) * len(ds.lon.values)) as bar:
        for lat in ds.lat.values:
            for lon in ds.lon.values:
                heat = np.zeros(len_date+n)
                i = 0
                while i < len_date:
                    value = ds.greater.sel(time=ds.time.values[i:i + n + 1], lat=lat, lon=lon)
                    if sum(value.values) == n:
                        heat[i] = 1
                        i+=3
                        bar(3) # Update the progress bar
                    else:
                        i+=1
                        bar() # Update the progress bar
                out.heatwave.loc[dict(lat=lat, lon=lon)] = heat
    return out

def heatwave_opt2_with_anomaly(ds:xr.Dataset, temperature:xr.Dataset, percent:xr.Dataset, n:int):
    # Cria uma cópia booleana com False por padrão
    ds = ds.astype(int)
    out = ds.copy()

    # Initialize temporal variables as boolean arrays
    out['heatwave'] = xr.zeros_like(ds.greater, dtype=bool)
    out['anomaly_tmax'] = xr.zeros_like(ds.greater, dtype=bool)
    out['anomaly_tmin'] = xr.zeros_like(ds.greater, dtype=bool)

    # Percorre cada data do dataset original
    len_date = len(ds.time[:-n+1])-1
    with alive_bar(len_date * len(ds.lat.values) * len(ds.lon.values)) as bar:
        for lat in ds.lat.values:
            for lon in ds.lon.values:
                heat = np.zeros(len_date+n)
                anom_tmax = np.zeros(len_date+n)
                anom_tmin = np.zeros(len_date+n)
                i = 0
                while i < len_date:
                    value = ds.greater.sel(time=ds.time.values[i:i + n + 1], lat=lat, lon=lon)
                    if sum(value.values) == n:
                        heat[i] = 1
                        anom_tmax[i] = temperature.tmax.sel(time=ds.time.values[i:i + n + 1], lat=lat, lon=lon).sum() - percent.tmax.sel(time=ds.time.values[i:i + n + 1], lat=lat, lon=lon).sum()
                        anom_tmax[i] = temperature.tmax.sel(time=ds.time.values[i:i + n + 1], lat=lat, lon=lon).sum() - percent.tmax.sel(time=ds.time.values[i:i + n + 1], lat=lat, lon=lon).sum()
                        i+=3
                        bar(3) # Update the progress bar
                    else:
                        i+=1
                        bar() # Update the progress bar
                out.heatwave.loc[dict(lat=lat, lon=lon)] = heat
                out.anomaly.loc[dict(lat=lat, lon=lon)] = heat
    return out