Active learning of a global surrogate model

Driver:

ActiveLearning

Download script:

active_surrogate_training.py

The target of the study is to train a surrogate model of a vectorial function. We consider an active training loop. That is, in each iteration, training data is generated by evaluating the vectorial function at the point of maximal prediction uncertainty.

As an example problem we consider the frequency dependent transmission function of a Fabry-Pérot etalon (angular frequency \(\omega\)). The etalon consists of a resonator of length \(l\) formed by a pair of mirrors with reflectivities \(R_1\) and \(R_2\). The propagation losses inside the resonator are quantified by the intensity-loss coefficient \(\alpha=0.05\).

The transmission function is given as (see also Wikipedia entry on Fabry-Pérot etalon)

\[A_\text{trans}(\omega, R_1, R_2, l) = \frac {(1-R_{1})(1-R_{2})e^{-\alpha l }}{\left({1-{\sqrt {R_{1}R_{2}}}e^{-\alpha l }}\right)^{2}+4{\sqrt {R_{1}R_{2}}}e^{-\alpha l }\sin^{2}(\phi )}.\]

The round-trip phase shift of the light field inside the resonator is \(\phi(\omega, l) = \omega \cdot l\).

We consider the case that we wish to learn the vectorial mapping from etalon parameters \((R_1, R_2, l)\) to transmission spectra

\[\begin{split}\mathbf{f}(R_1, R_2, l) = \begin{bmatrix} A_\text{trans}(\omega_1, R_1, R_2, l) \\ A_\text{trans}(\omega_2, R_1, R_2, l) \\ \vdots \\ A_\text{trans}(\omega_{50}, R_1, R_2, l) \end{bmatrix}\end{split}\]

with \(\omega_k = 2\pi k/50\).

It is possible to learn this vectorial mapping using multi-output a Gaussian process or a multi-output Bayesian neural network. Here, we present the approach to learn instead the 4d scalar function \(A_\text{trans}(\omega, R_1, R_2, l)\) using a single-output Bayesian neural network. This has the advantage that the vector entries are not learned independently, but that correlations between similar frequencies are taken into account. Moreover, after training one can get predictions for arbitrarily fine omega scans.

  1import sys,os
  2import numpy as np
  3import time
  4import matplotlib.pyplot as plt
  5
  6from jcmoptimizer import Client, Study, Obseravtion
  7client = Client()
  8
  9
 10def Atrans(
 11        R1: float, R2: float, l: float, alpha: float, omega: float
 12) -> np.ndarray:
 13    """Transmission through the etalon
 14    Args:
 15       R1: Reflectivity of first mirror
 16       R2: Reflectivity of second mirror
 17       l: resonator length 
 18       alpha: Intensity loss coefficient
 19       omega: Angular frequency of light
 20    """
 21    loss = np.exp(-alpha*l)
 22    R = np.sqrt(R1*R2)
 23    
 24    out = (1 - R1)*(1 - R2)*loss 
 25    out /= (1 - R*loss)**2 + 4*R*loss*np.sin(omega*l)**2
 26    return out
 27
 28# Definition of the parameter domain
 29design_space = [
 30    {'name': 'R1', 'type': 'continuous', 'domain': (0.1, 0.7)}, 
 31    {'name': 'R2', 'type': 'continuous', 'domain': (0.1, 0.7)}, 
 32    {'name': 'l', 'type': 'continuous', 'domain': (0.5, 1.0)}, 
 33]
 34
 35# Definition of the fixed environment variable alpha and the
 36# the scan variable omega
 37environment = [
 38    {'name': 'alpha', 'type': 'fixed', 'domain': 0.05},
 39    {'name': 'omega', 'type': 'variable', 'domain': (0, 2*np.pi)},
 40]
 41#The omega-scan defining the transmission spectra
 42omegas = np.linspace(0, 2*np.pi, 50)
 43
 44# Creation of the study object with study_id 'active_surrogate_training'
 45study = client.create_study(
 46    design_space=design_space,
 47    environment=environment,
 48    driver="ActiveLearning",
 49    study_name="Active learning of a global surrogate model",
 50    study_id="active_surrogate_training"
 51)
 52
 53study.configure(
 54    max_iter = 50,
 55    surrogates=[
 56        # We use a neural network with 4 hidden layers of 200 neurons each
 57        # to learn the scalar function Atrans(R1, R2, l, omega)
 58        dict(
 59            type="NN", name="Atrans", output_dim=1,
 60            hidden_layers_arch=[200, 200, 200, 200],
 61            num_NNs=60,
 62            optimization_step_max=-1,
 63            trainer=dict(
 64                type="full_data_trainer",
 65                num_epochs=1000,
 66                num_expel_NNs=30
 67            )
 68        )
 69    ],
 70    variables=[
 71        # The variable defines a scan of the surrogate prediction over all omega values
 72        dict(
 73            type="Scan",
 74            name="omega_scan",
 75            input_surrogate="Atrans",
 76            output_dim=len(omegas),
 77            scan_parameters=["omega"],
 78            scan_values=omegas[:, None].tolist(),
 79        ),
 80        # The variable defines the average transmission of the omega-scan
 81        dict(type="LinearCombination", name="average", inputs=["omega_scan"]),
 82    ],
 83    objectives=[
 84        # The objective is to evaluate the model function at maximal uncertainty of
 85        # the average transmission.
 86        dict(
 87            type="Explorer",
 88            name="objective",
 89            variable="average",
 90        )
 91    ]
 92)
 93
 94# Evaluation of the black-box function for specified design parameters
 95def evaluate(study: Study, R1: float, R2: float, l: float, alpha: float) -> Observation:
 96    time.sleep(2) # make objective expensive
 97    observation = study.new_observation()
 98    for omega in omegas:
 99        observation.add(
100            Atrans(R1, R2, l, alpha, omega),
101            environment_value=[omega],
102            model_name="Atrans",
103        )
104    return observation
105
106# Run the training loop
107study.set_evaluator(evaluate)
108study.run()
109
110
111
112study.configure(
113    surrogates=[
114        # For making more accurate predictions, we train the network on
115        # all data for 1500 epochs.
116        dict(
117            type="NN", name="Atrans", output_dim=1,
118            hidden_layers_arch=[200, 200, 200, 200],
119            num_NNs=60,
120            trainer=dict(
121                type="full_data_trainer",
122                num_epochs=1500,
123                num_expel_NNs=30
124            )
125        ),
126    ],
127)
128
129# Get prediction and anayltic values on a finer resolved omega-scan
130omegas_fine = np.linspace(0, 2 * np.pi, 150)
131
132# To test the worst-case prediction, we get a suggestion corresponding to
133# a sample with largest uncertainty
134s = study.get_suggestion()
135study.clear_suggestion(s.id)
136
137plt.figure(figsize=(10, 5))
138for R1, R2, l in [
139    (s.kwargs["R1"], s.kwargs["R2"], s.kwargs["l"]),
140    (0.1, 0.1, 0.5),
141    (0.1, 0.7, 0.75),
142    (0.7, 0.7, 1.0),
143]:
144    prediction = study.driver.predict(
145        points=[[R1, R2, l, omega] for omega in omegas_fine],
146        object_type="surrogate",
147        name="Atrans",
148    )
149    mean = np.array(prediction["mean"]).squeeze()
150    std = np.sqrt(np.array(prediction["variance"])).squeeze()
151    p = plt.plot(omegas_fine,  mean)
152    plt.fill_between(
153        omegas_fine, mean - std, mean + std, alpha=0.2, color=p[0].get_color(),
154    )
155    plt.plot(
156        omegas_fine,
157        [Atrans(R1, R2, l, 0.05, omega) for omega in omegas_fine],
158        "--",
159        color=p[0].get_color(),
160    )
161
162plt.xlabel("Angular frequency")
163plt.ylabel("Transmission")
164plt.grid()
165plt.savefig("etalon_predictions.svg", transparent=True)
166
167client.shutdown_server()
Prediction of etalon transmission

The figure shows for different parameters \(R_1, R_2\) and \(l\) the predicted transmission function (solid lines, shading indicates uncertainty of prediction) in comparison to the analytical transmission value (dashed lines). The blue line corresponds to the prediction with the largest average uncertainty. The other lines correspond to the etalon parameters \(R_1 = 0.1, R_2 = 0.1, l = 0.5\) (orange), \(R_1 = 0.1, R_2 = 0.7, l = 0.75\) (green) and \(R_1 = 0.7, R_2 = 0.7, l = 1.0\) (red). Considering the small number of 50 data points, the agreement between prediction and analytical value is very good.