sofm 0.1.4__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- sofm-0.1.4/LICENSE +21 -0
- sofm-0.1.4/PKG-INFO +130 -0
- sofm-0.1.4/README.md +101 -0
- sofm-0.1.4/pyproject.toml +38 -0
- sofm-0.1.4/setup.cfg +4 -0
- sofm-0.1.4/src/scov.py +309 -0
- sofm-0.1.4/src/sofm.egg-info/PKG-INFO +130 -0
- sofm-0.1.4/src/sofm.egg-info/SOURCES.txt +11 -0
- sofm-0.1.4/src/sofm.egg-info/dependency_links.txt +1 -0
- sofm-0.1.4/src/sofm.egg-info/requires.txt +14 -0
- sofm-0.1.4/src/sofm.egg-info/top_level.txt +3 -0
- sofm-0.1.4/src/sofm.py +117 -0
- sofm-0.1.4/src/utils.py +1544 -0
sofm-0.1.4/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 daniel-s-cunha
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
sofm-0.1.4/PKG-INFO
ADDED
|
@@ -0,0 +1,130 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: sofm
|
|
3
|
+
Version: 0.1.4
|
|
4
|
+
Summary: Spatially Orthogonal Factor Model (SOFM)
|
|
5
|
+
Author-email: Dan Cunha <dcunha@bu.edu>
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/daniel-s-cunha/SOFM
|
|
8
|
+
Classifier: Programming Language :: Python :: 3
|
|
9
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
10
|
+
Classifier: Operating System :: OS Independent
|
|
11
|
+
Requires-Python: >=3.8
|
|
12
|
+
Description-Content-Type: text/markdown
|
|
13
|
+
License-File: LICENSE
|
|
14
|
+
Requires-Dist: joblib>=1.3.0
|
|
15
|
+
Requires-Dist: matplotlib>=3.7.0
|
|
16
|
+
Requires-Dist: numpy>=1.24.0
|
|
17
|
+
Requires-Dist: pandas>=2.0.0
|
|
18
|
+
Requires-Dist: scikit-learn>=1.3.0
|
|
19
|
+
Requires-Dist: scikit-sparse>=0.4.12
|
|
20
|
+
Requires-Dist: scipy>=1.10.0
|
|
21
|
+
Requires-Dist: seaborn>=0.12.0
|
|
22
|
+
Requires-Dist: sympy>=1.12.0
|
|
23
|
+
Requires-Dist: threadpoolctl>=3.2.0
|
|
24
|
+
Requires-Dist: torch>=2.0.0
|
|
25
|
+
Requires-Dist: tqdm>=4.65.0
|
|
26
|
+
Requires-Dist: xarray>=2023.5.0
|
|
27
|
+
Requires-Dist: xeofs>=2.0.0
|
|
28
|
+
Dynamic: license-file
|
|
29
|
+
|
|
30
|
+
# SOFM
|
|
31
|
+
Spatially orthogonal factor models
|
|
32
|
+
|
|
33
|
+
## Installation
|
|
34
|
+
|
|
35
|
+
Please install `scikit-sparse` via conda before installing this package:
|
|
36
|
+
|
|
37
|
+
```bash
|
|
38
|
+
conda install -c conda-forge scikit-sparse
|
|
39
|
+
pip install sofm
|
|
40
|
+
```
|
|
41
|
+
|
|
42
|
+
## Data
|
|
43
|
+
Please use the `spatialLIBD_to_netcdf.R` script to download the DLPFC data from R and save as a netcdf file to be imported into xarray.
|
|
44
|
+
|
|
45
|
+
## Code example
|
|
46
|
+
```
|
|
47
|
+
from sofm import SOFM
|
|
48
|
+
|
|
49
|
+
# 1. Load netcdf into xarray
|
|
50
|
+
# Please see `spatialLIBD_to_netcdf.R` for directions on formatting the netcdf file.
|
|
51
|
+
da = xr.open_dataset("/data_directory/spatial_transcriptomic_data.nc", engine='netcdf4')['logcounts']
|
|
52
|
+
|
|
53
|
+
da = da - da.mean(dim='spot') #SOFM assumes the mean structure has been subtracted
|
|
54
|
+
da = da.rename({ #SOFM assumes `da` has `location` dimension indexed by `lat,lon` coordinates
|
|
55
|
+
'spot': 'location',
|
|
56
|
+
'array_col': 'lon',
|
|
57
|
+
'array_row': 'lat'
|
|
58
|
+
})
|
|
59
|
+
da = da.set_index(location=['lat', 'lon'])
|
|
60
|
+
|
|
61
|
+
# 2. Initialize SOFM model
|
|
62
|
+
sofm_model = SOFM(
|
|
63
|
+
data=da,
|
|
64
|
+
n_components=5,
|
|
65
|
+
n_cores=-1 #please set number of cores
|
|
66
|
+
)
|
|
67
|
+
|
|
68
|
+
# 3. Fit model
|
|
69
|
+
sofm_model.fit()
|
|
70
|
+
|
|
71
|
+
# 4. Visualize spatial loadings
|
|
72
|
+
fig = sofm_model.plot_loadings(robust=True)
|
|
73
|
+
fig.show()
|
|
74
|
+
|
|
75
|
+
# 5. Analyze latent factors
|
|
76
|
+
latent_factors = sofm_model.Ez_
|
|
77
|
+
```
|
|
78
|
+
|
|
79
|
+
If helpful, try out the code using a synthetic dataset,
|
|
80
|
+
```
|
|
81
|
+
import numpy as np
|
|
82
|
+
import xarray as xr
|
|
83
|
+
from sofm import SOFM
|
|
84
|
+
|
|
85
|
+
# 1. Generate Synthetic Data
|
|
86
|
+
# Create a 30x30 spatial grid
|
|
87
|
+
lats, lons = np.meshgrid(np.linspace(-5, 5, 30), np.linspace(-5, 5, 30))
|
|
88
|
+
lat_flat = lats.flatten()
|
|
89
|
+
lon_flat = lons.flatten()
|
|
90
|
+
n_locations = len(lat_flat)
|
|
91
|
+
n_features = 20
|
|
92
|
+
|
|
93
|
+
factor1 = np.sin(lat_flat) + np.cos(lon_flat)
|
|
94
|
+
factor2 = np.exp(-(lat_flat**2 + lon_flat**2) / 2)
|
|
95
|
+
factor3 = lat_flat + lon_flat
|
|
96
|
+
Z = np.column_stack([factor1, factor2, factor3])
|
|
97
|
+
|
|
98
|
+
np.random.seed(42)
|
|
99
|
+
W = np.random.randn(3, n_features) # Synthetic loadings
|
|
100
|
+
data_matrix = Z @ W + np.random.randn(n_locations, n_features) * 0.5
|
|
101
|
+
|
|
102
|
+
da = xr.DataArray(
|
|
103
|
+
data_matrix,
|
|
104
|
+
dims=['location', 'feature'],
|
|
105
|
+
coords={
|
|
106
|
+
'lat': ('location', lat_flat),
|
|
107
|
+
'lon': ('location', lon_flat),
|
|
108
|
+
'feature': np.arange(n_features)
|
|
109
|
+
}
|
|
110
|
+
)
|
|
111
|
+
da = da.set_index(location=['lat', 'lon'])
|
|
112
|
+
da = da - da.mean(dim='location')
|
|
113
|
+
|
|
114
|
+
# 2. Initialize SOFM model
|
|
115
|
+
sofm_model = SOFM(
|
|
116
|
+
data=da,
|
|
117
|
+
n_components=3,
|
|
118
|
+
n_cores=-1 #please set number of cores
|
|
119
|
+
)
|
|
120
|
+
|
|
121
|
+
# 3. Fit model
|
|
122
|
+
sofm_model.fit()
|
|
123
|
+
|
|
124
|
+
# 4. Visualize spatial loadings
|
|
125
|
+
fig = sofm_model.plot_loadings(robust=True)
|
|
126
|
+
fig.show()
|
|
127
|
+
|
|
128
|
+
# 5. Analyze latent factors
|
|
129
|
+
latent_factors = sofm_model.Ez_
|
|
130
|
+
```
|
sofm-0.1.4/README.md
ADDED
|
@@ -0,0 +1,101 @@
|
|
|
1
|
+
# SOFM
|
|
2
|
+
Spatially orthogonal factor models
|
|
3
|
+
|
|
4
|
+
## Installation
|
|
5
|
+
|
|
6
|
+
Please install `scikit-sparse` via conda before installing this package:
|
|
7
|
+
|
|
8
|
+
```bash
|
|
9
|
+
conda install -c conda-forge scikit-sparse
|
|
10
|
+
pip install sofm
|
|
11
|
+
```
|
|
12
|
+
|
|
13
|
+
## Data
|
|
14
|
+
Please use the `spatialLIBD_to_netcdf.R` script to download the DLPFC data from R and save as a netcdf file to be imported into xarray.
|
|
15
|
+
|
|
16
|
+
## Code example
|
|
17
|
+
```
|
|
18
|
+
from sofm import SOFM
|
|
19
|
+
|
|
20
|
+
# 1. Load netcdf into xarray
|
|
21
|
+
# Please see `spatialLIBD_to_netcdf.R` for directions on formatting the netcdf file.
|
|
22
|
+
da = xr.open_dataset("/data_directory/spatial_transcriptomic_data.nc", engine='netcdf4')['logcounts']
|
|
23
|
+
|
|
24
|
+
da = da - da.mean(dim='spot') #SOFM assumes the mean structure has been subtracted
|
|
25
|
+
da = da.rename({ #SOFM assumes `da` has `location` dimension indexed by `lat,lon` coordinates
|
|
26
|
+
'spot': 'location',
|
|
27
|
+
'array_col': 'lon',
|
|
28
|
+
'array_row': 'lat'
|
|
29
|
+
})
|
|
30
|
+
da = da.set_index(location=['lat', 'lon'])
|
|
31
|
+
|
|
32
|
+
# 2. Initialize SOFM model
|
|
33
|
+
sofm_model = SOFM(
|
|
34
|
+
data=da,
|
|
35
|
+
n_components=5,
|
|
36
|
+
n_cores=-1 #please set number of cores
|
|
37
|
+
)
|
|
38
|
+
|
|
39
|
+
# 3. Fit model
|
|
40
|
+
sofm_model.fit()
|
|
41
|
+
|
|
42
|
+
# 4. Visualize spatial loadings
|
|
43
|
+
fig = sofm_model.plot_loadings(robust=True)
|
|
44
|
+
fig.show()
|
|
45
|
+
|
|
46
|
+
# 5. Analyze latent factors
|
|
47
|
+
latent_factors = sofm_model.Ez_
|
|
48
|
+
```
|
|
49
|
+
|
|
50
|
+
If helpful, try out the code using a synthetic dataset,
|
|
51
|
+
```
|
|
52
|
+
import numpy as np
|
|
53
|
+
import xarray as xr
|
|
54
|
+
from sofm import SOFM
|
|
55
|
+
|
|
56
|
+
# 1. Generate Synthetic Data
|
|
57
|
+
# Create a 30x30 spatial grid
|
|
58
|
+
lats, lons = np.meshgrid(np.linspace(-5, 5, 30), np.linspace(-5, 5, 30))
|
|
59
|
+
lat_flat = lats.flatten()
|
|
60
|
+
lon_flat = lons.flatten()
|
|
61
|
+
n_locations = len(lat_flat)
|
|
62
|
+
n_features = 20
|
|
63
|
+
|
|
64
|
+
factor1 = np.sin(lat_flat) + np.cos(lon_flat)
|
|
65
|
+
factor2 = np.exp(-(lat_flat**2 + lon_flat**2) / 2)
|
|
66
|
+
factor3 = lat_flat + lon_flat
|
|
67
|
+
Z = np.column_stack([factor1, factor2, factor3])
|
|
68
|
+
|
|
69
|
+
np.random.seed(42)
|
|
70
|
+
W = np.random.randn(3, n_features) # Synthetic loadings
|
|
71
|
+
data_matrix = Z @ W + np.random.randn(n_locations, n_features) * 0.5
|
|
72
|
+
|
|
73
|
+
da = xr.DataArray(
|
|
74
|
+
data_matrix,
|
|
75
|
+
dims=['location', 'feature'],
|
|
76
|
+
coords={
|
|
77
|
+
'lat': ('location', lat_flat),
|
|
78
|
+
'lon': ('location', lon_flat),
|
|
79
|
+
'feature': np.arange(n_features)
|
|
80
|
+
}
|
|
81
|
+
)
|
|
82
|
+
da = da.set_index(location=['lat', 'lon'])
|
|
83
|
+
da = da - da.mean(dim='location')
|
|
84
|
+
|
|
85
|
+
# 2. Initialize SOFM model
|
|
86
|
+
sofm_model = SOFM(
|
|
87
|
+
data=da,
|
|
88
|
+
n_components=3,
|
|
89
|
+
n_cores=-1 #please set number of cores
|
|
90
|
+
)
|
|
91
|
+
|
|
92
|
+
# 3. Fit model
|
|
93
|
+
sofm_model.fit()
|
|
94
|
+
|
|
95
|
+
# 4. Visualize spatial loadings
|
|
96
|
+
fig = sofm_model.plot_loadings(robust=True)
|
|
97
|
+
fig.show()
|
|
98
|
+
|
|
99
|
+
# 5. Analyze latent factors
|
|
100
|
+
latent_factors = sofm_model.Ez_
|
|
101
|
+
```
|
|
@@ -0,0 +1,38 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=61.0.0", "wheel"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "sofm"
|
|
7
|
+
version = "0.1.4"
|
|
8
|
+
description = "Spatially Orthogonal Factor Model (SOFM)"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
authors = [
|
|
11
|
+
{ name = "Dan Cunha", email = "dcunha@bu.edu" }
|
|
12
|
+
]
|
|
13
|
+
license = { text = "MIT" }
|
|
14
|
+
classifiers = [
|
|
15
|
+
"Programming Language :: Python :: 3",
|
|
16
|
+
"License :: OSI Approved :: MIT License",
|
|
17
|
+
"Operating System :: OS Independent",
|
|
18
|
+
]
|
|
19
|
+
requires-python = ">=3.8"
|
|
20
|
+
dependencies = [
|
|
21
|
+
"joblib>=1.3.0",
|
|
22
|
+
"matplotlib>=3.7.0",
|
|
23
|
+
"numpy>=1.24.0",
|
|
24
|
+
"pandas>=2.0.0",
|
|
25
|
+
"scikit-learn>=1.3.0",
|
|
26
|
+
"scikit-sparse>=0.4.12",
|
|
27
|
+
"scipy>=1.10.0",
|
|
28
|
+
"seaborn>=0.12.0",
|
|
29
|
+
"sympy>=1.12.0",
|
|
30
|
+
"threadpoolctl>=3.2.0",
|
|
31
|
+
"torch>=2.0.0",
|
|
32
|
+
"tqdm>=4.65.0",
|
|
33
|
+
"xarray>=2023.5.0",
|
|
34
|
+
"xeofs>=2.0.0"
|
|
35
|
+
]
|
|
36
|
+
|
|
37
|
+
[project.urls]
|
|
38
|
+
Homepage = "https://github.com/daniel-s-cunha/SOFM"
|
sofm-0.1.4/setup.cfg
ADDED
sofm-0.1.4/src/scov.py
ADDED
|
@@ -0,0 +1,309 @@
|
|
|
1
|
+
from matplotlib.patches import Ellipse
|
|
2
|
+
import matplotlib.cm as cm
|
|
3
|
+
import time
|
|
4
|
+
from joblib import Parallel, delayed
|
|
5
|
+
import itertools
|
|
6
|
+
import matplotlib.pyplot as plt
|
|
7
|
+
import xarray as xr
|
|
8
|
+
import numpy as np
|
|
9
|
+
from numpy.linalg import svd
|
|
10
|
+
import torch
|
|
11
|
+
import torch.nn.functional as F
|
|
12
|
+
import torch.optim as optim
|
|
13
|
+
import pandas as pd
|
|
14
|
+
from sklearn.utils.extmath import randomized_svd
|
|
15
|
+
# import rioxarray as rxr
|
|
16
|
+
# from rasterio import CRS
|
|
17
|
+
import scipy.spatial
|
|
18
|
+
from scipy.spatial import cKDTree
|
|
19
|
+
from scipy.linalg import solve
|
|
20
|
+
import scipy.sparse as sp
|
|
21
|
+
from scipy.interpolate import BSpline
|
|
22
|
+
from sklearn.cluster import KMeans
|
|
23
|
+
from sksparse.cholmod import cholesky
|
|
24
|
+
#
|
|
25
|
+
import utils
|
|
26
|
+
from tqdm import tqdm
|
|
27
|
+
from threadpoolctl import threadpool_limits
|
|
28
|
+
|
|
29
|
+
class SpatialCovariance:
|
|
30
|
+
|
|
31
|
+
def __init__(self, data, nonstationary=True, n_components = 1, max_lag=30, block_sz = 5, n_blocks = 20, n_cores = -1):
|
|
32
|
+
self.data = self._standardize_input(data).compute()
|
|
33
|
+
self.nonstationary = nonstationary
|
|
34
|
+
self.max_lag = max_lag
|
|
35
|
+
self.block_sz = block_sz
|
|
36
|
+
self.n_blocks = n_blocks
|
|
37
|
+
self.n_components = n_components
|
|
38
|
+
self.n_cores = n_cores
|
|
39
|
+
#
|
|
40
|
+
self.holdout_ = None
|
|
41
|
+
self.holdout_centers_ = None
|
|
42
|
+
self.holdout_ids_ = None
|
|
43
|
+
self.spatcov_ = None
|
|
44
|
+
self.alpha_lat_ = None
|
|
45
|
+
self.alpha_lon_ = None
|
|
46
|
+
self.alpha_rot_ = None
|
|
47
|
+
self.t_u_ = None
|
|
48
|
+
self.t_v_ = None
|
|
49
|
+
self.lam_lat_ = None
|
|
50
|
+
self.lam_lon_ = None
|
|
51
|
+
self.sill_ = None
|
|
52
|
+
self.ls1_ = None
|
|
53
|
+
self.ls2_ = None
|
|
54
|
+
#
|
|
55
|
+
self.generate_holdout()
|
|
56
|
+
|
|
57
|
+
def _standardize_input(self, data):
|
|
58
|
+
if not isinstance(data, (xr.DataArray, xr.Dataset)):
|
|
59
|
+
raise TypeError("Input must be an xarray DataArray or Dataset.")
|
|
60
|
+
|
|
61
|
+
if 'location' in data.dims and not isinstance(data.indexes.get('location'), pd.MultiIndex):
|
|
62
|
+
if 'lat' in data.coords and 'lon' in data.coords:
|
|
63
|
+
return data.set_index(location=['lat', 'lon'])
|
|
64
|
+
|
|
65
|
+
if 'lat' in data.dims and 'lon' in data.dims:
|
|
66
|
+
return data.stack(location=['lat', 'lon'])
|
|
67
|
+
|
|
68
|
+
if 'location' in data.dims and isinstance(data.indexes.get('location'), pd.MultiIndex):
|
|
69
|
+
return data
|
|
70
|
+
|
|
71
|
+
raise ValueError("Data must contain 'lat' and 'lon' coordinates or dimensions.")
|
|
72
|
+
|
|
73
|
+
def generate_holdout(self):
|
|
74
|
+
#
|
|
75
|
+
self.holdout_, self.holdout_ids_, self.holdout_centers_ = utils._create_mask(self.data,n_blocks=self.n_blocks,block_sz=self.block_sz)
|
|
76
|
+
#
|
|
77
|
+
#return self
|
|
78
|
+
|
|
79
|
+
def plot_holdout(self):
|
|
80
|
+
M_da = xr.DataArray(
|
|
81
|
+
data=(self.holdout_!=0).numpy().astype(float),
|
|
82
|
+
dims=('location'),
|
|
83
|
+
coords={
|
|
84
|
+
'location': self.data.location
|
|
85
|
+
}
|
|
86
|
+
)
|
|
87
|
+
M_da = M_da.unstack()
|
|
88
|
+
M_da = M_da.sortby(['lon','lat'])
|
|
89
|
+
M_da.plot(add_colorbar=False)
|
|
90
|
+
plt.show()
|
|
91
|
+
|
|
92
|
+
def plot_nonstationary_spatcov(self, invert=True, robust=False):
|
|
93
|
+
if not self.nonstationary:
|
|
94
|
+
raise ValueError(
|
|
95
|
+
"This plot can only be used when nonstationary=True. "
|
|
96
|
+
"The current model is initialized as stationary."
|
|
97
|
+
)
|
|
98
|
+
|
|
99
|
+
if not hasattr(self, 'lam_lat_'):
|
|
100
|
+
raise ValueError("plot_nonstationary_spatcov can only be used after .fit() is completed.")
|
|
101
|
+
|
|
102
|
+
# --- Data Preparation ---
|
|
103
|
+
def prep_da(data):
|
|
104
|
+
da = xr.DataArray(data=data, dims=('location'), coords={'location': self.data.location})
|
|
105
|
+
return da.unstack().sortby('lon', 'lat')
|
|
106
|
+
|
|
107
|
+
W_da = prep_da(self.lam_lat_)
|
|
108
|
+
W_da2 = prep_da(self.lam_lon_)
|
|
109
|
+
W_da3 = prep_da(self.rot_)
|
|
110
|
+
|
|
111
|
+
ar = np.maximum(self.lam_lat_, self.lam_lon_) / (np.minimum(self.lam_lat_, self.lam_lon_) + 1e-8)
|
|
112
|
+
W_da4 = prep_da(ar)
|
|
113
|
+
|
|
114
|
+
# --- Plotting Configuration ---
|
|
115
|
+
title_size = 24
|
|
116
|
+
label_size = 14
|
|
117
|
+
fig, axes = plt.subplots(1, 2, figsize=(12, 6)) # Slightly larger fig for larger text
|
|
118
|
+
|
|
119
|
+
# Helper to apply labels consistently
|
|
120
|
+
def format_ax(ax, title, is_first=False):
|
|
121
|
+
ax.set_title(title, fontsize=title_size, pad=10)
|
|
122
|
+
ax.set_xlabel('', fontsize=label_size)
|
|
123
|
+
ax.set_ylabel('' if is_first else '', fontsize=label_size)
|
|
124
|
+
# Increase tick label size
|
|
125
|
+
ax.tick_params(labelsize=10)
|
|
126
|
+
|
|
127
|
+
# 1. Latitudinal
|
|
128
|
+
W_da.plot(ax=axes[0], cmap=cm.plasma, robust=robust, #vmin=3, vmax=5,
|
|
129
|
+
add_colorbar=True, cbar_kwargs={'format': '%.1f'})
|
|
130
|
+
format_ax(axes[0], 'Latitudinal length scale', is_first=True)
|
|
131
|
+
|
|
132
|
+
# 2. Longitudinal
|
|
133
|
+
W_da2.plot(ax=axes[1], cmap=cm.plasma, robust=robust,#vmin=3, vmax=5,
|
|
134
|
+
add_colorbar=True, cbar_kwargs={'format': '%.1f'})
|
|
135
|
+
format_ax(axes[1], 'Longitudinal length scale')
|
|
136
|
+
|
|
137
|
+
# # 3. Rotation
|
|
138
|
+
# W_da3.plot(ax=axes[2], cmap=cm.plasma, vmin=0, robust=True,
|
|
139
|
+
# add_colorbar=True, cbar_kwargs={'format': '%.2f'})
|
|
140
|
+
# format_ax(axes[2], 'Rotation')
|
|
141
|
+
|
|
142
|
+
# # 4. Anisotropy
|
|
143
|
+
# W_da4.plot(ax=axes[2], cmap=cm.plasma, add_colorbar=True, robust=True,
|
|
144
|
+
# cbar_kwargs={'format': '%.1f'})
|
|
145
|
+
# format_ax(axes[2], 'Anisotropy ratio')
|
|
146
|
+
if invert:
|
|
147
|
+
for ax in axes:
|
|
148
|
+
ax.invert_yaxis()
|
|
149
|
+
plt.tight_layout()
|
|
150
|
+
return fig
|
|
151
|
+
|
|
152
|
+
def _evaluate_ls(self, ls1, ls2, rot, init_phi):
|
|
153
|
+
Sigma = utils._compute_spat_cov_rs(
|
|
154
|
+
self.data,
|
|
155
|
+
phi=init_phi,
|
|
156
|
+
length_scale=ls1,
|
|
157
|
+
length_scale2=ls2,
|
|
158
|
+
rot = rot,
|
|
159
|
+
max_lag=self.max_lag
|
|
160
|
+
)
|
|
161
|
+
U, L, Ez, sigma2, loss, loss_tot = utils._cv_spatPCA(
|
|
162
|
+
self.data,
|
|
163
|
+
Sigma,
|
|
164
|
+
self.holdout_,
|
|
165
|
+
k=self.n_components,
|
|
166
|
+
phi=init_phi
|
|
167
|
+
)
|
|
168
|
+
if self.nonstationary:
|
|
169
|
+
return (ls1, ls2, rot, loss)
|
|
170
|
+
else:
|
|
171
|
+
return (ls1, ls2, rot, loss_tot)
|
|
172
|
+
|
|
173
|
+
def _evaluate_phi(self, phi):
|
|
174
|
+
Sigma = utils._compute_spat_cov_rs(
|
|
175
|
+
self.data,
|
|
176
|
+
phi=1,
|
|
177
|
+
length_scale=1,
|
|
178
|
+
max_lag=self.max_lag
|
|
179
|
+
)
|
|
180
|
+
#
|
|
181
|
+
U, L, Ez, sigma2, loss, loss_tot = utils._cv_spatPCA(
|
|
182
|
+
self.data,
|
|
183
|
+
phi*Sigma,
|
|
184
|
+
self.holdout_,
|
|
185
|
+
k=self.n_components,
|
|
186
|
+
phi=phi
|
|
187
|
+
)
|
|
188
|
+
|
|
189
|
+
return (phi, loss_tot)
|
|
190
|
+
|
|
191
|
+
def _evaluate_phi_post(self, phi):
|
|
192
|
+
Sigma = self.spatcov_
|
|
193
|
+
#
|
|
194
|
+
U, L, Ez, sigma2, loss, loss_tot = utils._cv_spatPCA(
|
|
195
|
+
self.data,
|
|
196
|
+
phi*Sigma,
|
|
197
|
+
self.holdout_,
|
|
198
|
+
k=self.n_components,
|
|
199
|
+
phi=phi
|
|
200
|
+
)
|
|
201
|
+
|
|
202
|
+
return (phi, loss_tot)
|
|
203
|
+
|
|
204
|
+
def fit(self, lss=[1,3,5,7], phis=[1e3,5e3,1e4,5e4], rots=[0, 1*np.pi/16, 2*np.pi/16, 3*np.pi/16]):
|
|
205
|
+
#
|
|
206
|
+
if self.nonstationary:
|
|
207
|
+
self._fit_nonstationary(lss, phis, rots)
|
|
208
|
+
else:
|
|
209
|
+
self._fit_stationary(lss, phis)
|
|
210
|
+
|
|
211
|
+
def _fit_nonstationary(self, lss, phis, rots):
|
|
212
|
+
#
|
|
213
|
+
#
|
|
214
|
+
#fit best phi for fixed ls
|
|
215
|
+
#
|
|
216
|
+
if len(phis)>1:
|
|
217
|
+
results = list(tqdm(
|
|
218
|
+
Parallel(n_jobs=self.n_cores, return_as="generator")(
|
|
219
|
+
delayed(self._evaluate_phi)(phi)
|
|
220
|
+
for phi in phis
|
|
221
|
+
),
|
|
222
|
+
total = len(phis),
|
|
223
|
+
desc = 'Validating sill'
|
|
224
|
+
))
|
|
225
|
+
#
|
|
226
|
+
self.sill_, _ = min(results, key=lambda x: x[1])
|
|
227
|
+
else:
|
|
228
|
+
self.sill_ = phis[0]
|
|
229
|
+
#
|
|
230
|
+
print(f'The prior sill estimate is {self.sill_}')
|
|
231
|
+
n_centers = len(self.holdout_ids_)
|
|
232
|
+
#
|
|
233
|
+
# k_1d = np.sqrt(n_centers)
|
|
234
|
+
# opt_knots = int(max(3, np.round(k_1d - 2)))
|
|
235
|
+
#
|
|
236
|
+
#fit best lss
|
|
237
|
+
ls_combinations = list(itertools.product(lss, lss, rots))
|
|
238
|
+
init_phi = self.sill_
|
|
239
|
+
#
|
|
240
|
+
results = list(tqdm(
|
|
241
|
+
Parallel(n_jobs=self.n_cores, return_as='generator')(
|
|
242
|
+
delayed(self._evaluate_ls)(ls1, ls2, rot, init_phi)
|
|
243
|
+
for ls1, ls2, rot in ls_combinations
|
|
244
|
+
),
|
|
245
|
+
total=len(ls_combinations),
|
|
246
|
+
desc='Validating length scales'
|
|
247
|
+
))
|
|
248
|
+
|
|
249
|
+
res_dict = {(res[0], res[1], res[2]): res[3] for res in results}
|
|
250
|
+
loss_df = pd.DataFrame(res_dict)
|
|
251
|
+
|
|
252
|
+
loss_df.index = self.holdout_ids_
|
|
253
|
+
|
|
254
|
+
minimizer = loss_df.idxmin(axis=1).rename('ls_hat')
|
|
255
|
+
|
|
256
|
+
centers_df = pd.DataFrame(
|
|
257
|
+
self.holdout_centers_,
|
|
258
|
+
index=self.holdout_ids_,
|
|
259
|
+
columns=['mean_lat', 'mean_lon']
|
|
260
|
+
)
|
|
261
|
+
|
|
262
|
+
final_results = pd.concat([centers_df, minimizer], axis=1).dropna()
|
|
263
|
+
|
|
264
|
+
self.minimizer_ = np.column_stack([
|
|
265
|
+
np.array(final_results['ls_hat'].tolist()),
|
|
266
|
+
final_results['mean_lat'].values,
|
|
267
|
+
final_results['mean_lon'].values
|
|
268
|
+
])
|
|
269
|
+
print("Fitting spline to the optimal length scales...")
|
|
270
|
+
self.alpha_lat_, self.alpha_lon_, self.alpha_rot_, self.t_u_, self.t_v_ = utils._fit_spline(self.data,self.minimizer_)
|
|
271
|
+
#fit with variance=1 so you only have to run it once and can instead adjust it by multiplying phi
|
|
272
|
+
self.spatcov_, self.lam_lat_, self.lam_lon_, self.rot_ = utils._construct_nonstat_cov(self.data.lat.values, self.data.lon.values, self.alpha_lat_, self.alpha_lon_, self.alpha_rot_, self.t_u_, self.t_v_, variance=1, max_lag=self.max_lag)
|
|
273
|
+
#
|
|
274
|
+
#
|
|
275
|
+
self.spatcov_ = self.sill_ * self.spatcov_
|
|
276
|
+
|
|
277
|
+
def _fit_stationary(self, lss, phis):
|
|
278
|
+
#
|
|
279
|
+
#fit best ls for fixed init_phi
|
|
280
|
+
#
|
|
281
|
+
ls_combinations = list(itertools.product(lss, lss))
|
|
282
|
+
init_phi = 1e4
|
|
283
|
+
|
|
284
|
+
results = list(tqdm(
|
|
285
|
+
Parallel(n_jobs=self.n_cores, return_as='generator')(
|
|
286
|
+
delayed(self._evaluate_ls)(ls1, ls2, [0], init_phi)
|
|
287
|
+
for ls1, ls2 in ls_combinations
|
|
288
|
+
),
|
|
289
|
+
total = len(ls_combinations),
|
|
290
|
+
desc = 'Validating length scales'
|
|
291
|
+
))
|
|
292
|
+
|
|
293
|
+
self.ls1_, self.ls2_, _, _ = min(results, key=lambda x: x[2])
|
|
294
|
+
self.spatcov_ = utils._compute_spat_cov_rs(self.data, 1, self.ls1_, self.ls2_, max_lag=self.max_lag)
|
|
295
|
+
#
|
|
296
|
+
#fit best phi for fixed ls
|
|
297
|
+
#
|
|
298
|
+
results = list(tqdm(
|
|
299
|
+
Parallel(n_jobs=self.n_cores,return_as='generator')(
|
|
300
|
+
delayed(self._evaluate_phi)(phi)
|
|
301
|
+
for phi in phis
|
|
302
|
+
),
|
|
303
|
+
total = len(phis),
|
|
304
|
+
desc = 'Validating sill'
|
|
305
|
+
))
|
|
306
|
+
|
|
307
|
+
self.sill_, _ = min(results, key=lambda x: x[1])
|
|
308
|
+
self.spatcov_ = self.sill_ * self.spatcov_
|
|
309
|
+
|