EnVarPractical/utilities.py at develop · benjaminmenetrier/EnVarPractical · 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
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Tools required for the EnVar exercises:
- define_correlation_sqrt(n, R): return the square-root of a correlation matrix of size (n,n), with a correlation cut-off length-scale R
- generate_ensemble: returns an ensemble of size nens, randomizing the correlation matrix of size (n,n), with a correlation cut-off length-scale R
- plot_vector(X, title): plot a vector (or a set of vectors) X with the provided title
- plot_matrix(A, title): plot matrix A with the provided title
- plot_matrices(A1, title1, A2, title2): plot matrix A1 and A2 with the provided titles title1 and title2, respectively.

Programming: Benjamin Menetrier ([email protected])

Licensing: this code is distributed under the APACHE 2 license
Copyright (c) 2023-2024 Meteorologisk Institutt
"""
###############################################################################
# Imports
###############################################################################
import matplotlib.pyplot as plt
import numpy as np

###############################################################################
# define_correlation_sqrt
###############################################################################
def define_correlation_sqrt(n, R):
    U = np.zeros((n, n))
    for i in range(n):
        for j in range(n):
            dist = np.abs(i-j)
            if (dist>n/2):
                dist = n-dist
            dist = dist/R
            if dist <= 0.5:
                U[i,j] = 1.0-(2.0*dist)
    x = np.zeros((n, 1))
    x[0,0] = 1.0
    x = np.matmul(np.transpose(U), x)
    x = np.matmul(U, x)
    U = U/np.sqrt(x[0,0])
    return U

###############################################################################
# generate_ensemble
###############################################################################
def generate_ensemble(n, Ul, nens):
    # Generate Gaussian deviates
    X = np.random.normal(0.0, 1.0, (n, nens))

    # Apply localization square-root
    XX = np.zeros((n, nens))
    for j in range(nens):
        XX[:,j] = np.matmul(Ul, X[:,j])

    return XX

###############################################################################
# plot_vector
###############################################################################
def plot_vector(X, title):
    n, nens = np.shape(X)
    fig, ax = plt.subplots(figsize=(14,7))
    ax.set_xlim([0,n-1])
    ax.set_title(title)
    ax.axhline(y=0, color='r')
    ax.plot(X, 'k')
    plt.show()

###############################################################################
# plot_matrix
###############################################################################
def plot_matrix(A, title):
    fig = plt.figure()
    plt.set_cmap('jet')
    ax = fig.add_subplot(111)
    ax.set_title(title)
    cax = ax.matshow(A, vmin=-0.2, vmax=1.2)
    fig.colorbar(cax)
    plt.show()

###############################################################################
# plot_matrices
###############################################################################
def plot_matrices(A1, title1, A2, title2):
    fig = plt.figure()
    plt.set_cmap('jet')
    ax = fig.add_subplot(121)
    ax.set_title(title1)
    cax = ax.matshow(A1, vmin=-0.2, vmax=1.2)
    ax = fig.add_subplot(122)
    ax.set_title(title2)
    cax = ax.matshow(A2, vmin=-0.2, vmax=1.2)
    plt.show()