This notebook is available at arjunsavel/cortecs

Fitting with a neural network

In the Quickstart tutorial, we worked through how to use the PCA-based approach implemented in cortecs. Now, let’s walk through how to use the neural network-based approach.

Setting up the objects

import numpy as np

import cortecs
from cortecs.opac.opac import *
from cortecs.fit.fit import *
from cortecs.eval.eval import *

We’ll be using the same Opac object as in the Quickstart.

T_filename = "temperatures.npy"
P_filename = "pressures.npy"
wl_filename = "wavelengths.npy"

cross_sec_filename = "absorb_coeffs_C2H4.npy"

load_kwargs = {
    "T_filename": T_filename,
    "P_filename": P_filename,
    "wl_filename": wl_filename,
}
opac_obj = Opac(cross_sec_filename, loader="platon", load_kwargs=load_kwargs)

This time, we specify that we’d like to fit the object with the neural_net method.

fitter = Fitter(opac_obj, method="neural_net")
fitter

<cortecs.fit.fit.Fitter at 0x29217e950>

We won’t actually fit the entire opacity, though—that would take about 15 hours on a laptop. Instead, let’s just fit a single one using a lower-level API.

res = cortecs.fit.fit_neural_net.fit_neural_net(
    fitter.opac.cross_section[:, :, -2],
    fitter.opac.T,
    fitter.opac.P,
    None,
    n_layers=3,
    n_neurons=8,
    activation="sigmoid",
    learn_rate=0.04,
    loss="mean_squared_error",
    epochs=4000,
    verbose=0,
    sequential_model=None,
    plot=True,
)

This code shouldn’t take more than 15 seconds or so to run.

The above plot tells us how the model performs as it iterates on its weights.

history, neural_network = res

P_unraveled = unravel_data(fitter.opac.P, fitter.opac.T, None, tileboth=True)
T_unraveled = unravel_data(fitter.opac.T, fitter.opac.P, None, tileboth=False)
input_array = np.column_stack([T_unraveled, P_unraveled])

npres = len(fitter.opac.P)
ntemp = len(fitter.opac.T)

predictions = neural_network.predict(input_array)
plt.imshow(
    100
    * (predictions.reshape(ntemp, npres) - fitter.opac.cross_section[:, :, -1])
    / predictions.reshape(ntemp, npres)
)
plt.colorbar()

13/13 [==============================] - 0s 571us/step

<matplotlib.colorbar.Colorbar at 0x29e981660>

Awesome! Looks like we have accuracy at better than ~5%. We can tune the architecture of the neural network to achieve greater accuracy, but keep in mind that a larger network means more memory — that is, less efficient compression.

We can save the weights for this individual neural network with the built-in API.

save_neural_net(neural_network, "test_nn.pkl")

Great! We’ve saved our network to the disk. Let’s make sure that evaluating this on the fly (without the neural network API — just the weights and the biases) works out.

First, we load in the weights and the biases.

with open("test_nn.pkl", "rb") as f:
    all_weights, all_biases = pickle.load(f)

Next, we call the eval_neural_net method to evaluate the neural net with those weights and biases.

n_layers = len(all_weights)
eval_neural_net(100, 1e-4, n_layers, all_weights, all_biases)

Array([-49.47523], dtype=float32)

Great! And let’s check that this matches the prediction we made earlier with the same neural net.

predictions[0]

array([-49.475227], dtype=float32)

Finally, let’s do a speed test.

%%timeit
eval_neural_net(2000, 1e-4, n_layers, all_weights, all_biases)

214 µs ± 1.96 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

Again, this is slower than simply accessing the array. This is a trade-off you’ll have to make depending on your architecture (CPU vs. GPU), desired accuracy, and desired memory performance.

This tutorial has walked through how to use the neural network-based approach to fitting opacity data. We did so using a single neural network architecture — i.e., we only used n_layers=3, n_neurons=8, etc. These model “hyperparameters” in principle have an optimal value for a given dataset, given constraints on model size. We will explore this idea further in the Optimizing Fits tutorial.