pyteda 0.2.0__tar.gz → 0.2.5__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.
Files changed (38) hide show
  1. pyteda-0.2.5/PKG-INFO +196 -0
  2. pyteda-0.2.5/README.md +174 -0
  3. {pyteda-0.2.0 → pyteda-0.2.5}/pyproject.toml +1 -1
  4. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda/analysis/analysis_enkf_modified_cholesky.py +3 -0
  5. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda/analysis/analysis_lenkf.py +5 -4
  6. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda/analysis/analysis_letkf.py +6 -4
  7. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda/models/__init__.py +1 -0
  8. pyteda-0.2.5/pyteda/models/quasi_geostrophic.py +178 -0
  9. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda/simulation/simulation_core.py +46 -19
  10. pyteda-0.2.5/pyteda.egg-info/PKG-INFO +196 -0
  11. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda.egg-info/SOURCES.txt +1 -0
  12. pyteda-0.2.0/PKG-INFO +0 -346
  13. pyteda-0.2.0/README.md +0 -324
  14. pyteda-0.2.0/pyteda.egg-info/PKG-INFO +0 -346
  15. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda/__init__.py +0 -0
  16. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda/analysis/__init__.py +0 -0
  17. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda/analysis/analysis_core.py +0 -0
  18. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda/analysis/analysis_enkf.py +0 -0
  19. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda/analysis/analysis_enkf_bloc.py +0 -0
  20. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda/analysis/analysis_enkf_cholesky.py +0 -0
  21. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda/analysis/analysis_enkf_naive.py +0 -0
  22. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda/analysis/analysis_enkf_shrinkage_precision.py +0 -0
  23. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda/analysis/analysis_ensrf.py +0 -0
  24. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda/analysis/analysis_etkf.py +0 -0
  25. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda/analysis/analysis_factory.py +0 -0
  26. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda/analysis/loader.py +0 -0
  27. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda/analysis/registry.py +0 -0
  28. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda/background/__init__.py +0 -0
  29. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda/background/background_core.py +0 -0
  30. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda/models/lorenz96.py +0 -0
  31. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda/models/model.py +0 -0
  32. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda/observation/__init__.py +0 -0
  33. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda/observation/observation_core.py +0 -0
  34. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda/simulation/__init__.py +0 -0
  35. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda.egg-info/dependency_links.txt +0 -0
  36. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda.egg-info/requires.txt +0 -0
  37. {pyteda-0.2.0 → pyteda-0.2.5}/pyteda.egg-info/top_level.txt +0 -0
  38. {pyteda-0.2.0 → pyteda-0.2.5}/setup.cfg +0 -0
pyteda-0.2.5/PKG-INFO ADDED
@@ -0,0 +1,196 @@
1
+ Metadata-Version: 2.4
2
+ Name: pyteda
3
+ Version: 0.2.5
4
+ Summary: TEDA: Toy models and ensemble-based Data Assimilation methods for education and experimentation
5
+ Author-email: "Elías D. Niño-Ruiz" <elias.d.nino@gmail.com>
6
+ License: MIT
7
+ Keywords: data assimilation,EnKF,education,Lorenz96,Python
8
+ Classifier: Development Status :: 3 - Alpha
9
+ Classifier: Intended Audience :: Education
10
+ Classifier: Topic :: Scientific/Engineering :: Atmospheric Science
11
+ Classifier: License :: OSI Approved :: MIT License
12
+ Classifier: Programming Language :: Python :: 3
13
+ Classifier: Programming Language :: Python :: 3.7
14
+ Classifier: Programming Language :: Python :: 3.8
15
+ Classifier: Programming Language :: Python :: 3.9
16
+ Requires-Python: >=3.7
17
+ Description-Content-Type: text/markdown
18
+ Requires-Dist: numpy
19
+ Requires-Dist: scipy
20
+ Requires-Dist: matplotlib
21
+ Requires-Dist: scikit-learn
22
+
23
+ # 🧠 TEDA – Toolbox for Ensemble Data Assimilation
24
+
25
+ We are thrilled to present TEDA, a cutting-edge Python toolbox designed to facilitate the teaching and learning of ensemble-based data assimilation (DA). TEDA caters to the needs of educators and learners alike, offering a comprehensive platform to explore the captivating world of meteorological anomalies, climate change, and DA methodologies.
26
+
27
+ TEDA stands out as an exceptional resource in the landscape of DA software, as it prioritizes the educational experience. While operational software tends to focus on practical applications, TEDA empowers undergraduate and graduate students by providing an intuitive platform to grasp the intricacies of ensemble-based DA concepts and foster enthusiasm for scientific exploration.
28
+
29
+ Equipped with a diverse range of features and functionalities, TEDA enriches the learning process by offering powerful visualization tools. Students can delve into error statistics, model errors, observational errors, error distributions, and the dynamic evolution of errors. These visualizations provide multiple perspectives on numerical results, enabling students to develop a deep understanding of ensemble-based DA principles.
30
+
31
+ TEDA offers numerous ensemble-based DA methods, allowing students to explore an extensive range of techniques. From the stochastic ensemble Kalman filter (EnKF) to dual EnKF formulations, EnKF via Cholesky decomposition, EnKF based on modified Cholesky decomposition, EnKF based on B-localization, and beyond, TEDA equips learners with a comprehensive toolkit to study and compare various DA methodologies.
32
+
33
+ TEDA's foundation is rooted in the Object-Oriented Programming (OOP) paradigm, enabling effortless integration of new techniques and models. This ensures that TEDA remains at the forefront of the rapidly evolving field of ensemble-based DA, allowing educators and researchers to incorporate the latest advancements into their teaching and experimentation.
34
+
35
+ Embark on a captivating journey of discovery and knowledge with TEDA as it unveils the secrets of data assimilation. Simulate various DA scenarios, experiment with different model configurations, and witness firsthand the transformative impact of ensemble-based DA methods on forecast accuracy and data analysis.
36
+
37
+ Unlock your true potential in the realm of data assimilation with TEDA. Whether you are an enthusiastic student or a dedicated educator, TEDA stands as your ultimate companion in unraveling the complexities of ensemble-based DA.
38
+
39
+ Keywords: Data Assimilation, Ensemble Kalman Filter, Education, Python.
40
+
41
+ ### 🧪 Toy Models for Data Assimilation
42
+
43
+ To enhance the learning experience, **TEDA** is built around **toy models** that exhibit chaotic behavior under specific configurations. These models are ideal for experimenting with and validating various Data Assimilation (DA) techniques in a hands-on, controlled setting.
44
+
45
+ Currently, TEDA includes the **Lorenz 96 model** (40 variables), a well-known benchmark in DA research.
46
+
47
+ > 📌 *Support for other classic chaotic systems like the Duffing equation and Lorenz 63 may be added in future versions.*
48
+
49
+ ---
50
+
51
+ ## 📦 Installation
52
+
53
+ Aquí tienes una versión mejor redactada y estructurada para la sección de instalación y uso rápido:
54
+
55
+ ---
56
+
57
+ ## 🚀 Installation and Quick Start
58
+
59
+ ### ✅ From PyPI
60
+
61
+ To install the latest stable version:
62
+
63
+ ```bash
64
+ pip install pyteda
65
+ ```
66
+
67
+ ### 🧪 From Source (for development)
68
+
69
+ Clone the repository and install in editable mode:
70
+
71
+ ```bash
72
+ git clone https://github.com/your-username/pyteda.git
73
+ cd pyteda
74
+ pip install -e .
75
+ ```
76
+
77
+ ### 🛠️ Recommended Setup: Virtual Environment
78
+
79
+ We suggest using a virtual environment to manage dependencies:
80
+
81
+ ```bash
82
+ python -m venv venv
83
+ source venv/bin/activate # On Windows: venv\Scripts\activate
84
+ ```
85
+
86
+ Install all required dependencies:
87
+
88
+ ```bash
89
+ pip install -r requirements.txt
90
+ ```
91
+
92
+ ---
93
+
94
+ ## 📁 Example Usage
95
+
96
+ You can run a complete simulation workflow using:
97
+
98
+ ```bash
99
+ python examples/basic_template.py
100
+ ```
101
+
102
+ Or explore the file directly: [`examples/basic_template.py`](examples/basic_template.py)
103
+
104
+ ---
105
+
106
+ ### ▶️ Run the basic example
107
+
108
+ Once installed, you can run the full example with:
109
+
110
+ ```bash
111
+ python -m examples.basic_template
112
+ ```
113
+
114
+ This will initialize the Lorenz96 model, run an ensemble-based data assimilation simulation, and display the log-relative errors for both the background and analysis states.
115
+
116
+ ---
117
+
118
+ ### Supported Methods
119
+
120
+ The `pyteda` package supports a wide range of ensemble-based data assimilation methods including EnKF, LETKF, ETKF, and shrinkage-based filters. All methods are accessible via the `AnalysisFactory`.
121
+
122
+ 📄 See [Supported Methods](docs/index.md) for a full list and references.
123
+
124
+ ---
125
+
126
+ ## 🚀 Quickstart
127
+
128
+ Here's a minimal example using the **LETKF** method with the built-in Lorenz96 model:
129
+
130
+ ```python
131
+ from pyteda.simulation import Simulation
132
+ from pyteda.models import Lorenz96
133
+ from pyteda.background import Background
134
+ from pyteda.observation import Observation
135
+ from pyteda.analysis.analysis_factory import AnalysisFactory
136
+
137
+ model = Lorenz96()
138
+ background = Background(model, ensemble_size=20)
139
+ observation = Observation(m=32, std_obs=0.01)
140
+
141
+ analysis = AnalysisFactory("letkf", model=model).create_analysis()
142
+
143
+ params = {'obs_freq': 0.1, 'obs_times': 10, 'inf_fact': 1.04}
144
+ sim = Simulation(model, background, analysis, observation, params=params)
145
+ sim.run()
146
+
147
+ errb, erra = sim.get_errors()
148
+ ```
149
+
150
+ 📈 For visualization and full working examples, check the [`examples/`](examples/) folder.
151
+
152
+ ---
153
+
154
+ ## 🎓 Citation
155
+
156
+ If you use **TEDA** in your teaching or research, please cite:
157
+
158
+ **Nino-Ruiz, E.D., Valbuena, S.R. (2022)**
159
+ *TEDA: A Computational Toolbox for Teaching Ensemble Based Data Assimilation*.
160
+ [📖 ICCS 2022, Springer](https://doi.org/10.1007/978-3-031-08760-8_60)
161
+
162
+ ```bibtex
163
+ @inproceedings{nino2022teda,
164
+ title={TEDA: A Computational Toolbox for Teaching Ensemble Based Data Assimilation},
165
+ author={Nino-Ruiz, Elias D. and Valbuena, Sebastian Racedo},
166
+ booktitle={Computational Science – ICCS 2022},
167
+ series={Lecture Notes in Computer Science},
168
+ volume={13353},
169
+ pages={787--801},
170
+ year={2022},
171
+ publisher={Springer, Cham},
172
+ doi={10.1007/978-3-031-08760-8_60}
173
+ }
174
+ ```
175
+
176
+ ---
177
+
178
+ ## 📚 More Resources
179
+
180
+ * 🧪 Explore [examples/](examples/) to try different filters and models.
181
+ * 📖 Full documentation available in the `docs/` folder.
182
+ * ➕ Want to add your own method? TEDA supports [custom DA methods via registry](docs/custom_methods.md).
183
+
184
+ ---
185
+
186
+ [![PyPI version](https://badge.fury.io/py/pyteda.svg)](https://pypi.org/project/pyteda/)
187
+ [![Cite TEDA](https://img.shields.io/badge/Cite-TEDA-blue.svg)](https://doi.org/10.1007/978-3-031-08760-8_60)
188
+
189
+ ---
190
+
191
+ ## Developed by
192
+
193
+ * Elías D. Niño-Ruiz
194
+ * [https://enino84.github.io/](https://enino84.github.io/)
195
+ * elias.d.nino@gmail.com
196
+
pyteda-0.2.5/README.md ADDED
@@ -0,0 +1,174 @@
1
+ # 🧠 TEDA – Toolbox for Ensemble Data Assimilation
2
+
3
+ We are thrilled to present TEDA, a cutting-edge Python toolbox designed to facilitate the teaching and learning of ensemble-based data assimilation (DA). TEDA caters to the needs of educators and learners alike, offering a comprehensive platform to explore the captivating world of meteorological anomalies, climate change, and DA methodologies.
4
+
5
+ TEDA stands out as an exceptional resource in the landscape of DA software, as it prioritizes the educational experience. While operational software tends to focus on practical applications, TEDA empowers undergraduate and graduate students by providing an intuitive platform to grasp the intricacies of ensemble-based DA concepts and foster enthusiasm for scientific exploration.
6
+
7
+ Equipped with a diverse range of features and functionalities, TEDA enriches the learning process by offering powerful visualization tools. Students can delve into error statistics, model errors, observational errors, error distributions, and the dynamic evolution of errors. These visualizations provide multiple perspectives on numerical results, enabling students to develop a deep understanding of ensemble-based DA principles.
8
+
9
+ TEDA offers numerous ensemble-based DA methods, allowing students to explore an extensive range of techniques. From the stochastic ensemble Kalman filter (EnKF) to dual EnKF formulations, EnKF via Cholesky decomposition, EnKF based on modified Cholesky decomposition, EnKF based on B-localization, and beyond, TEDA equips learners with a comprehensive toolkit to study and compare various DA methodologies.
10
+
11
+ TEDA's foundation is rooted in the Object-Oriented Programming (OOP) paradigm, enabling effortless integration of new techniques and models. This ensures that TEDA remains at the forefront of the rapidly evolving field of ensemble-based DA, allowing educators and researchers to incorporate the latest advancements into their teaching and experimentation.
12
+
13
+ Embark on a captivating journey of discovery and knowledge with TEDA as it unveils the secrets of data assimilation. Simulate various DA scenarios, experiment with different model configurations, and witness firsthand the transformative impact of ensemble-based DA methods on forecast accuracy and data analysis.
14
+
15
+ Unlock your true potential in the realm of data assimilation with TEDA. Whether you are an enthusiastic student or a dedicated educator, TEDA stands as your ultimate companion in unraveling the complexities of ensemble-based DA.
16
+
17
+ Keywords: Data Assimilation, Ensemble Kalman Filter, Education, Python.
18
+
19
+ ### 🧪 Toy Models for Data Assimilation
20
+
21
+ To enhance the learning experience, **TEDA** is built around **toy models** that exhibit chaotic behavior under specific configurations. These models are ideal for experimenting with and validating various Data Assimilation (DA) techniques in a hands-on, controlled setting.
22
+
23
+ Currently, TEDA includes the **Lorenz 96 model** (40 variables), a well-known benchmark in DA research.
24
+
25
+ > 📌 *Support for other classic chaotic systems like the Duffing equation and Lorenz 63 may be added in future versions.*
26
+
27
+ ---
28
+
29
+ ## 📦 Installation
30
+
31
+ Aquí tienes una versión mejor redactada y estructurada para la sección de instalación y uso rápido:
32
+
33
+ ---
34
+
35
+ ## 🚀 Installation and Quick Start
36
+
37
+ ### ✅ From PyPI
38
+
39
+ To install the latest stable version:
40
+
41
+ ```bash
42
+ pip install pyteda
43
+ ```
44
+
45
+ ### 🧪 From Source (for development)
46
+
47
+ Clone the repository and install in editable mode:
48
+
49
+ ```bash
50
+ git clone https://github.com/your-username/pyteda.git
51
+ cd pyteda
52
+ pip install -e .
53
+ ```
54
+
55
+ ### 🛠️ Recommended Setup: Virtual Environment
56
+
57
+ We suggest using a virtual environment to manage dependencies:
58
+
59
+ ```bash
60
+ python -m venv venv
61
+ source venv/bin/activate # On Windows: venv\Scripts\activate
62
+ ```
63
+
64
+ Install all required dependencies:
65
+
66
+ ```bash
67
+ pip install -r requirements.txt
68
+ ```
69
+
70
+ ---
71
+
72
+ ## 📁 Example Usage
73
+
74
+ You can run a complete simulation workflow using:
75
+
76
+ ```bash
77
+ python examples/basic_template.py
78
+ ```
79
+
80
+ Or explore the file directly: [`examples/basic_template.py`](examples/basic_template.py)
81
+
82
+ ---
83
+
84
+ ### ▶️ Run the basic example
85
+
86
+ Once installed, you can run the full example with:
87
+
88
+ ```bash
89
+ python -m examples.basic_template
90
+ ```
91
+
92
+ This will initialize the Lorenz96 model, run an ensemble-based data assimilation simulation, and display the log-relative errors for both the background and analysis states.
93
+
94
+ ---
95
+
96
+ ### Supported Methods
97
+
98
+ The `pyteda` package supports a wide range of ensemble-based data assimilation methods including EnKF, LETKF, ETKF, and shrinkage-based filters. All methods are accessible via the `AnalysisFactory`.
99
+
100
+ 📄 See [Supported Methods](docs/index.md) for a full list and references.
101
+
102
+ ---
103
+
104
+ ## 🚀 Quickstart
105
+
106
+ Here's a minimal example using the **LETKF** method with the built-in Lorenz96 model:
107
+
108
+ ```python
109
+ from pyteda.simulation import Simulation
110
+ from pyteda.models import Lorenz96
111
+ from pyteda.background import Background
112
+ from pyteda.observation import Observation
113
+ from pyteda.analysis.analysis_factory import AnalysisFactory
114
+
115
+ model = Lorenz96()
116
+ background = Background(model, ensemble_size=20)
117
+ observation = Observation(m=32, std_obs=0.01)
118
+
119
+ analysis = AnalysisFactory("letkf", model=model).create_analysis()
120
+
121
+ params = {'obs_freq': 0.1, 'obs_times': 10, 'inf_fact': 1.04}
122
+ sim = Simulation(model, background, analysis, observation, params=params)
123
+ sim.run()
124
+
125
+ errb, erra = sim.get_errors()
126
+ ```
127
+
128
+ 📈 For visualization and full working examples, check the [`examples/`](examples/) folder.
129
+
130
+ ---
131
+
132
+ ## 🎓 Citation
133
+
134
+ If you use **TEDA** in your teaching or research, please cite:
135
+
136
+ **Nino-Ruiz, E.D., Valbuena, S.R. (2022)**
137
+ *TEDA: A Computational Toolbox for Teaching Ensemble Based Data Assimilation*.
138
+ [📖 ICCS 2022, Springer](https://doi.org/10.1007/978-3-031-08760-8_60)
139
+
140
+ ```bibtex
141
+ @inproceedings{nino2022teda,
142
+ title={TEDA: A Computational Toolbox for Teaching Ensemble Based Data Assimilation},
143
+ author={Nino-Ruiz, Elias D. and Valbuena, Sebastian Racedo},
144
+ booktitle={Computational Science – ICCS 2022},
145
+ series={Lecture Notes in Computer Science},
146
+ volume={13353},
147
+ pages={787--801},
148
+ year={2022},
149
+ publisher={Springer, Cham},
150
+ doi={10.1007/978-3-031-08760-8_60}
151
+ }
152
+ ```
153
+
154
+ ---
155
+
156
+ ## 📚 More Resources
157
+
158
+ * 🧪 Explore [examples/](examples/) to try different filters and models.
159
+ * 📖 Full documentation available in the `docs/` folder.
160
+ * ➕ Want to add your own method? TEDA supports [custom DA methods via registry](docs/custom_methods.md).
161
+
162
+ ---
163
+
164
+ [![PyPI version](https://badge.fury.io/py/pyteda.svg)](https://pypi.org/project/pyteda/)
165
+ [![Cite TEDA](https://img.shields.io/badge/Cite-TEDA-blue.svg)](https://doi.org/10.1007/978-3-031-08760-8_60)
166
+
167
+ ---
168
+
169
+ ## Developed by
170
+
171
+ * Elías D. Niño-Ruiz
172
+ * [https://enino84.github.io/](https://enino84.github.io/)
173
+ * elias.d.nino@gmail.com
174
+
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "pyteda"
3
- version = "0.2.0"
3
+ version = "0.2.5"
4
4
  description = "TEDA: Toy models and ensemble-based Data Assimilation methods for education and experimentation"
5
5
  readme = "README.md"
6
6
  requires-python = ">=3.7"
@@ -52,6 +52,9 @@ class AnalysisEnKFModifiedCholesky(Analysis):
52
52
  D[0, 0] = 1 / np.var(DX[0, :]) # We are estimating D^{-1}
53
53
  for i in range(1, n):
54
54
  ind_prede = self.model.get_pre(i, self.r)
55
+ if len(ind_prede) == 0:
56
+ D[i, i] = 1 / np.var(y)
57
+ continue
55
58
  y = DX[i, :]
56
59
  X = DX[ind_prede, :].T
57
60
  lr_fit = lr.fit(X, y)
@@ -24,7 +24,8 @@ class AnalysisLEnKF(Analysis):
24
24
 
25
25
  def local_analysis_LEnKF(self, Xb, H, R, y, n, N, i, r):
26
26
  # Subdomain decomposition
27
- si = [(i+j) % n for j in range(-r, r+1)]
27
+ si = si = self.model.get_ngb(i, r) #[(i+j) % n for j in range(-r, r+1)]
28
+ center_index = np.where(si == i)[0][0] # Index of the center element in the subdomain
28
29
  Xbi = Xb[:, si]
29
30
  Pbi = np.cov(Xbi.T)
30
31
  yz = np.zeros((n,))
@@ -49,7 +50,7 @@ class AnalysisLEnKF(Analysis):
49
50
  else:
50
51
  Xai = Xbi
51
52
 
52
- return Xai
53
+ return Xai, center_index # Return the analysis state and the index of the center element in the subdomain
53
54
 
54
55
  def perform_assimilation(self, background, observation):
55
56
  """
@@ -76,8 +77,8 @@ class AnalysisLEnKF(Analysis):
76
77
 
77
78
  Xa = np.zeros((ensemble_size, n)) # Local analysis for each model component i
78
79
  for i in range(0, n):
79
- Xai = self.local_analysis_LEnKF(Xb, H, R, y, n, ensemble_size, i, self.r)
80
- Xa[:, i] = Xai[:, self.r] #
80
+ Xai, center_index = self.local_analysis_LEnKF(Xb, H, R, y, n, ensemble_size, i, self.r)
81
+ Xa[:, i] = Xai[:, center_index]
81
82
 
82
83
  self.Xa = Xa.T
83
84
 
@@ -23,7 +23,8 @@ class AnalysisLETKF(Analysis):
23
23
 
24
24
  def local_analysis_LETKF(self, Xb, H, R, y, n, N, i, r):
25
25
  # Subdomain decomposition
26
- si = [(i+j) % n for j in range(-r, r+1)]
26
+ si = self.model.get_ngb(i, r) #[(i+j) % n for j in range(-r, r+1)]
27
+ center_index = np.where(si == i)[0][0] # Index of the center element in the subdomain
27
28
  Xbi = Xb[:, si].T
28
29
  xbi = np.mean(Xbi, axis=1).reshape(-1,1)
29
30
  DXi = Xbi - xbi
@@ -66,7 +67,7 @@ class AnalysisLETKF(Analysis):
66
67
  else:
67
68
  Xai = Xbi.T
68
69
 
69
- return Xai
70
+ return Xai, center_index
70
71
 
71
72
  def perform_assimilation(self, background, observation):
72
73
  """
@@ -93,8 +94,9 @@ class AnalysisLETKF(Analysis):
93
94
 
94
95
  Xa = np.zeros((ensemble_size, n)) # Local analysis for each model component i
95
96
  for i in range(0, n):
96
- Xai = self.local_analysis_LETKF(Xb, H, R, y, n, ensemble_size, i, self.r)
97
- Xa[:, i] = Xai[:, self.r]
97
+ Xai, center_index = self.local_analysis_LETKF(Xb, H, R, y, n, ensemble_size, i, self.r)
98
+ Xa[:, i] = Xai[:, center_index]
99
+ #Xa[:, i] = Xai[:, self.r]
98
100
 
99
101
  self.Xa = Xa.T
100
102
 
@@ -6,3 +6,4 @@ Created on Thu Oct 22 00:43:05 2020
6
6
  """
7
7
 
8
8
  from .lorenz96 import Lorenz96
9
+ from .quasi_geostrophic import QGModel
@@ -0,0 +1,178 @@
1
+ import numpy as np
2
+ from scipy.fft import fft2, ifft2, fftfreq
3
+ from .model import Model
4
+
5
+ class QGModel(Model):
6
+ """Barotropic QG model with periodic BCs and state vector [q, psi]."""
7
+
8
+ def __init__(self, N=64, L=1.0, dt=0.001, F=1600.0,
9
+ r=1e-2, rkb=1e-1, rkh=1e-5, rkh2=1e-8):
10
+ self.N = N
11
+ self.L = L
12
+ self.dt = dt
13
+ self.F = F
14
+ self.r = r
15
+ self.rkb = rkb
16
+ self.rkh = rkh
17
+ self.rkh2 = rkh2
18
+ self.n = 2 * N * N # state = [q, psi]
19
+ self._L = None
20
+
21
+ x = np.linspace(0, L, N, endpoint=False)
22
+ y = np.linspace(0, L, N, endpoint=False)
23
+ self.X, self.Y = np.meshgrid(x, y)
24
+ self.dx = self.dy = L / N
25
+
26
+ kx = fftfreq(N, d=self.dx) * 2 * np.pi
27
+ ky = fftfreq(N, d=self.dy) * 2 * np.pi
28
+ self.KX, self.KY = np.meshgrid(kx, ky)
29
+ self.K2 = self.KX**2 + self.KY**2
30
+ self.K2[0, 0] = 1e-10
31
+ self.K4 = self.K2**2
32
+
33
+ def laplacian(self, psi):
34
+ return ifft2(-self.K2 * fft2(psi)).real
35
+
36
+ def biharmonic(self, psi):
37
+ return ifft2(self.K4 * fft2(psi)).real
38
+
39
+ def arakawa_jacobian(self, psi, q):
40
+ def jp(f): return np.roll(f, -1, axis=0)
41
+ def jm(f): return np.roll(f, 1, axis=0)
42
+ def kp(f): return np.roll(f, -1, axis=1)
43
+ def km(f): return np.roll(f, 1, axis=1)
44
+
45
+ j1 = (jp(psi) - jm(psi)) * (kp(q) - km(q)) - (kp(psi) - km(psi)) * (jp(q) - jm(q))
46
+ j2 = kp(psi)*(jp(q)-q) - km(psi)*(q-jm(q)) - jp(psi)*(kp(q)-q) + jm(psi)*(q-km(q))
47
+ j3 = jp(psi)*(kp(q)-km(q)) - jm(psi)*(kp(q)-km(q)) - kp(psi)*(jp(q)-jm(q)) + km(psi)*(jp(q)-jm(q))
48
+ return (j1 + j2 + j3) / (12 * self.dx * self.dy)
49
+
50
+ def forcing(self):
51
+ return 2 * np.pi * np.sin(2 * np.pi * self.Y)
52
+
53
+ def propagate(self, x0, T, just_final_state=True):
54
+ q = x0[:self.N**2].reshape((self.N, self.N)).copy()
55
+ psi_hat = fft2(q) / (self.K2 + self.F)
56
+ psi_hat[0, 0] = 0.0
57
+ psi = ifft2(psi_hat).real
58
+
59
+ t = T[0]
60
+ t_end = T[-1]
61
+
62
+ while t < t_end:
63
+ t += self.dt
64
+ zeta = self.laplacian(psi)
65
+ q = zeta - self.F * psi
66
+ j = self.arakawa_jacobian(psi, q)
67
+ dqdt = -psi.copy()
68
+ dqdt += -self.r * j
69
+ dqdt += -self.rkb * zeta
70
+ dqdt += self.rkh * self.laplacian(zeta)
71
+ dqdt += -self.rkh2 * self.biharmonic(zeta)
72
+ dqdt += self.forcing()
73
+
74
+ q += self.dt * dqdt
75
+ psi_hat = fft2(q) / (self.K2 + self.F)
76
+ psi_hat[0, 0] = 0.0
77
+ psi = ifft2(psi_hat).real
78
+
79
+ if just_final_state:
80
+ return np.concatenate([q.flatten(), psi.flatten()])
81
+ return None
82
+
83
+ def get_initial_condition(self, seed=42, T=np.arange(0, 25, 0.01)):
84
+ np.random.seed(seed)
85
+ psi0 = 1e-4 * np.random.randn(self.N, self.N)
86
+ q0 = self.laplacian(psi0) - self.F * psi0
87
+ return self.propagate(np.concatenate([q0.flatten(), psi0.flatten()]), T)
88
+
89
+ def get_number_of_variables(self):
90
+ return self.n
91
+
92
+ def get_ngb(self, i, r, cross=False):
93
+ N = self.N
94
+ n_grid = N * N
95
+
96
+ var_index = 0 if i < n_grid else 1 # 0 = q, 1 = psi
97
+ i_local = i % n_grid
98
+ iy, ix = divmod(i_local, N)
99
+
100
+ neighbors = []
101
+ for dy in range(-r, r + 1):
102
+ y = (iy + dy) % N
103
+ for dx in range(-r, r + 1):
104
+ x = (ix + dx) % N
105
+ base = y * N + x
106
+
107
+ # Add same variable neighbor
108
+ neighbors.append(base + var_index * n_grid)
109
+
110
+ if cross:
111
+ # Add cross-variable neighbor
112
+ neighbors.append(base + (1 - var_index) * n_grid)
113
+
114
+ return np.array(sorted(set(neighbors)), dtype=int) # Remove duplicates if cross=True
115
+
116
+
117
+ def create_decorrelation_matrix(self, r, cross=False, cross_scale=1.0):
118
+ N = self.N
119
+ n = N * N
120
+ total_n = 2 * n
121
+ L = np.zeros((total_n, total_n))
122
+
123
+ for i in range(total_n):
124
+ iy1, ix1 = divmod(i % n, N)
125
+ var1 = 0 if i < n else 1 # 0: q, 1: psi
126
+
127
+ for j in range(i, total_n):
128
+ iy2, ix2 = divmod(j % n, N)
129
+ var2 = 0 if j < n else 1 # 0: q, 1: psi
130
+
131
+ dx = min(abs(ix1 - ix2), N - abs(ix1 - ix2))
132
+ dy = min(abs(iy1 - iy2), N - abs(iy1 - iy2))
133
+ dij2 = dx**2 + dy**2
134
+
135
+ if var1 == var2:
136
+ scale = 1.0
137
+ elif cross:
138
+ scale = cross_scale
139
+ else:
140
+ scale = 0.0
141
+
142
+ val = np.exp(-dij2 / (2 * r**2)) * scale
143
+ L[i, j] = L[j, i] = val
144
+
145
+ self._L = L
146
+
147
+
148
+ def get_decorrelation_matrix(self):
149
+ return self._L
150
+
151
+ def get_pre(self, i, r, cross=False):
152
+ N = self.N
153
+ n = N * N
154
+ total_n = 2 * n
155
+
156
+ if i < n:
157
+ offset = 0 # q
158
+ else:
159
+ offset = n # psi
160
+ i_base = i - n
161
+ i_base = i % n
162
+
163
+ iy, ix = divmod(i_base, N)
164
+ neighbors = []
165
+
166
+ for dy in range(-r, r + 1):
167
+ y = (iy + dy) % N
168
+ for dx in range(-r, r + 1):
169
+ x = (ix + dx) % N
170
+ j_base = y * N + x
171
+
172
+ for block_offset in [0, n]: # q and psi blocks
173
+ j = j_base + block_offset
174
+ if j < i:
175
+ if cross or block_offset == offset:
176
+ neighbors.append(j)
177
+
178
+ return np.array(neighbors, dtype=int)