stresearch · ktrapeznikov · May 24, 2024 · Jul 20, 2022 · Jul 20, 2022 · Jul 20, 2022
diff --git a/docs/index.md b/docs/index.md
@@ -17,7 +17,14 @@ Code repository and documentation for:
 ![](sections/overview/str.png)   ![](https://www.unsw.edu.au/content/dam/images/graphics/logos/unsw/unsw_0.png)
 
 - [STR](https://www.str.us) 
-- [University of New South Wales, Sydney](https://www.ccrc.unsw.edu.au/ccrc-team/academic-research/steven-sherwood)
+  -   Jim Gimlett
+  -   Gregory Bushey
+  -   Kirill Trapeznikov
+  -   Eric Jones
+- [University of New South Wales, Sydney](https://www.unsw.edu.au/)
+  -  [Steven Sherwood](https://www.unsw.edu.au/staff/steven-sherwood)
+  -  [Abhnil Prasad](https://www.unsw.edu.au/staff/abhnil-prasad)
+  -  David Fuchs 
 
 ## Sections
 - [Datasets](sections/datasets/README.md) - description of datasets used in surrogate training
@@ -29,7 +36,9 @@ Code repository and documentation for:
 - [MJO Analysis](sections/mjo_analysis/README.md) - analysis of MJO in the datasets
 - [Surrogate Bottleneck Analysis](sections/bottleneck/README.md) - analysis of bottleneck dimension of the surrogate model
 - [Linearization](sections/linearization/README.md) - analysis linearization of surrogate models to understand sensitivity to inputs
-
+- [Surrogate Architectures](sections/architectures/README.md) - comparison of different neural network architectures and configurations
+- [Fine Tuning](sections/finetuning/README.md) - comparison of fine-tuning an CAM4 trained model on SPCAM vs a model traiend from scratch
+- [Integration Journal](sections/integration/README.md) - details on integration including training setup, diagnostic plots, etc.
 
 ## Overview
 

diff --git a/docs/milestone_reports/milestone_report_5.pdf b/docs/milestone_reports/milestone_report_5.pdf
diff --git a/docs/milestone_reports/milestone_report_5_old.pdf b/docs/milestone_reports/milestone_report_5_old.pdf
diff --git a/docs/milestone_reports/milestone_report_6.pdf b/docs/milestone_reports/milestone_report_6.pdf
diff --git a/docs/sections/architectures/README.md b/docs/sections/architectures/README.md
@@ -0,0 +1,17 @@
+# Comparison of Different Surrogate Neural Network Architectures
+
+We compare skill of trained surrogates with different architectures. The models are designed to have approx. same number of parameters: `1.5 millions`.
+
+There are two variations:
+
+| Architecture | Description |  
+| -- | -- |  
+| **Non-spatial** | *models treat each grid column independently* |  
+ | fcn | fully connected network, 7 dense layers |  
+| conv1d_k | 1 dimensional conv net with dialation, z / levels = sequence dimension, variables = channel dimension, k = kernel size |  
+ | transformer |  transformer encoder model with z / level position encoding, z / levels = sequence dimension, variables = channel dimension |  
+| **Spatial** | *models can use information from neighbouring grids in making predictions* |  
+| *conv2d_k* | 2 dimensional seperable depthwise conv net, lat/lots = 2d spatial dimensions, variables stacked as channels, k = kernel size |  
+
+
+[![](skill_vs_nn_arch.png)](skill_vs_nn_arch.html)
diff --git a/docs/sections/architectures/skill_vs_nn_arch.html b/docs/sections/architectures/skill_vs_nn_arch.html
diff --git a/docs/sections/architectures/skill_vs_nn_arch.png b/docs/sections/architectures/skill_vs_nn_arch.png
diff --git a/docs/sections/finetuning/README.md b/docs/sections/finetuning/README.md
@@ -0,0 +1,18 @@
+# Finetuning Experiments
+
+Our strategy is to iteravly refine neural networks surrogates by training on them on progessively higher fidelity physics simulation data.
+
+To test our our fine-tuning strategy, we fine-tune a CAM4 trained NN surrogate on an increasing amount of SPCAM data. We compare fine-tuned model to a regular model trained from SCRATCH on the SPCAM data using the same amount data.
+
+The following plot shows the benefit of fine-tuning vs training from scratch.
+
+![](spcam-funetune.png)
+
+Training paramaters:
+- Total number of samples, N = 18e6, (in SPCAM training data)
+- Subsample factor, S = 2 ^ [3, 4, …. , 17]
+- For all, max_epochs: 200 / log(S)
+- For all, batch_size = min( 24 x 96 x 144,  N / S)
+- Fine-tune: learning rate 5e-5, no learning rate schedule
+- Train from scratch/random init: learning rage 1e-3, warm up schedule (ramp from 0 to lr during first 10% of epochs and then ramp down from lr to 0 during last 90%)
+- Best model over epochs taken w.r.t. to validation set
diff --git a/docs/sections/finetuning/spcam-funetune.png b/docs/sections/finetuning/spcam-funetune.png
diff --git a/docs/sections/integration/README.md b/docs/sections/integration/README.md
diff --git a/docs/sections/linearization/README.md b/docs/sections/linearization/README.md
@@ -2,12 +2,55 @@
 
 We compute linearization of the NN surrigate f(x) = y, to understand sensitivity to different inputs.
 
-Linearization of f(x`) = Ax` + b, can be computed by computing the gradiant grad f(x) around x`. Since f(x) is vector valued we compute a jacobian: A = J_x(f(x)). 
+Linearization of $f(x) = Ax + b$, can be computed by taking the gradiant  $grad~f(x)$ around $x$. Since f(x) is vector valued we compute a jacobian instead: $A = J_x(f(x))$. 
 
-To visualize how much each input perturbs every output, we l2 normalize each row of J. The following figure breaks up J by different output and input groups. This is computed around the mean input.
 
-## CAM4
+
+
+
+##  Jacobian of Temperature and Humidity Tendencies
+
+- We can compare empirical Jacobian of surrogate NN with reference CRM computations.
+- While the general structure is present in the surrogate, it still somewhat noisy. 
+- We plan to add Jacobian structure as an additional regularization during optimization to (hopefully) encourage stability during intergration into the climate model.
+> To compute the Jacobian, we sample 10000 data inputs from desired region, compute a jacobian around each sample using  NN autodifferentiation and then average. We found that this computation is more stable then computing Jacobian around a mean input.
+
+### Comparison to Reference CRM
+
+[*Click on figure for interactive visualization*](jacobian_proper_scale.html)
+
+Jacobian CAM4 FCN NN  |  Reference from CRM (Fig 8. Kuang 2012)
+:-------------------------:  |      :-------------------------:
+[![](jacobian_proper_scale.png)](jacobian_proper_scale.html) | [![](ref_jacobian.png)](ref_jacobian.png)
+
+
+Jacobian SPCAM FCN NN  |  Reference from CRM (Fig 8. Kuang 2012)
+:-------------------------:  |      :-------------------------:
+[![](jacobian_proper_scale_spcam.png)](jacobian_proper_scale_spcam.html) | [![](ref_jacobian.png)](ref_jacobian.png)
+
+<!-- [*By different lats*](jacobian_norm_cam4_by_lats.html) -->
+
+
+<!-- ### l2 Normalized Jacobian of All Inputs
+
+To visualize how much each input perturbs every output, we l2 normalize each row of J. The following figure breaks up J by different output and input groups. This is computed around the mean input.
 
 [*Click on figure for interactive visualization*](jacobian.html)
+[![](jacobian.png)](jacobian.html) -->
+
+
+### Architecture Comparison
+
+Jacobian CAM4 NN
+
+[![](jacobian_compare_cam4.png)](jacobian_compare_cam4.html)
+
+Jacobian SPCAM NN
+
+
+[![](jacobian_compare_spcam.png)](jacobian_compare_spcam.html)
+
+
+### Bottleneck
 
-[![](jacobian.png)](jacobian.html)
+[![](jacobian_bottleneck_cam4.png)](jacobian_bottleneck_cam4.html)
diff --git a/docs/sections/linearization/jacobian_bottleneck_cam4.html b/docs/sections/linearization/jacobian_bottleneck_cam4.html
diff --git a/docs/sections/linearization/jacobian_bottleneck_cam4.png b/docs/sections/linearization/jacobian_bottleneck_cam4.png
diff --git a/docs/sections/linearization/jacobian_compare_cam4.html b/docs/sections/linearization/jacobian_compare_cam4.html
diff --git a/docs/sections/linearization/jacobian_compare_cam4.png b/docs/sections/linearization/jacobian_compare_cam4.png
diff --git a/docs/sections/linearization/jacobian_compare_spcam.html b/docs/sections/linearization/jacobian_compare_spcam.html
diff --git a/docs/sections/linearization/jacobian_compare_spcam.png b/docs/sections/linearization/jacobian_compare_spcam.png
diff --git a/docs/sections/linearization/jacobian_proper_scale.html b/docs/sections/linearization/jacobian_proper_scale.html
diff --git a/docs/sections/linearization/jacobian_proper_scale.png b/docs/sections/linearization/jacobian_proper_scale.png
diff --git a/docs/sections/linearization/jacobian_proper_scale_spcam.html b/docs/sections/linearization/jacobian_proper_scale_spcam.html
diff --git a/docs/sections/linearization/jacobian_proper_scale_spcam.png b/docs/sections/linearization/jacobian_proper_scale_spcam.png
diff --git a/docs/sections/linearization/ref_jacobian.png b/docs/sections/linearization/ref_jacobian.png
diff --git a/docs/talks/.gitkeep b/docs/talks/.gitkeep
@@ -0,0 +1 @@
+