import one example into another to avoid lengthy code duplication

Examples/fit55_SpecularIntro/PolarizedSpinAsymmetry.py

+/G/sw/ba/Examples/varia/PolarizedSpinAsymmetry.py
\ No newline at end of file

Examples/fit55_SpecularIntro/PolarizedSpinAsymmetryFit.py

@@ -16,77 +16,7 @@ import matplotlib.pyplot as plt
 import bornagain as ba
 from bornagain import deg, angstrom
 
-# q-range on which the simulation and fitting are to be performed
-qmin = 0.05997
-qmax = 1.96
-
-# number of points on which the computed result is plotted
-scan_size = 1500
-
-# The SLD of the substrate is kept constant
-sldMao = (5.377e-06, 0)
-
-# constant to convert between magnetization and magnetic SLD
-RhoMconst = 2.910429812376859e-12
-
-####################################################################
-#                  Create Sample and Simulation                    #
-####################################################################
-
-
-def get_sample(params):
-    """
-    construct the sample with the given parameters
-    """
-    magnetizationMagnitude = params["rhoM_Mafo"]*1e-6/RhoMconst
-    angle = 0
-    magnetizationVector = ba.R3(
-        magnetizationMagnitude*numpy.sin(angle*deg),
-        magnetizationMagnitude*numpy.cos(angle*deg), 0)
-
-    mat_vacuum = ba.MaterialBySLD("Vacuum", 0, 0)
-    mat_layer = ba.MaterialBySLD("(Mg,Al,Fe)3O4", params["rho_Mafo"]*1e-6,
-                                 0, magnetizationVector)
-    mat_substrate = ba.MaterialBySLD("MgAl2O4", *sldMao)
-
-    ambient_layer = ba.Layer(mat_vacuum)
-    layer = ba.Layer(mat_layer, params["t_Mafo"]*angstrom)
-    substrate_layer = ba.Layer(mat_substrate)
-
-    r_Mafo = ba.LayerRoughness(params["r_Mafo"]*angstrom)
-    r_substrate = ba.LayerRoughness(params["r_Mao"]*angstrom)
-
-    multi_layer = ba.MultiLayer()
-    multi_layer.addLayer(ambient_layer)
-    multi_layer.addLayerWithTopRoughness(layer, r_Mafo)
-    multi_layer.addLayerWithTopRoughness(substrate_layer, r_substrate)
-
-    return multi_layer
-
-
-def get_simulation(q_axis, parameters, polarizer_dir, analyzer_dir):
-    """
-    Returns a simulation object.
-    Polarization, analyzer and resolution are set
-    from given parameters
-    """
-    simulation = ba.SpecularSimulation()
-    q_axis = q_axis + parameters["q_offset"]
-    scan = ba.QzScan(q_axis)
-
-    # unused, why was this define: dq = parameters["q_res"]*q_axis
-    n_sig = 4.0
-    n_samples = 25
-
-    distr = ba.RangedDistributionGaussian(n_samples, n_sig)
-    scan.setAbsoluteQResolution(distr, parameters["q_res"])
-
-    simulation.instrument().setPolFilters(polarizer_dir, analyzer_dir, 1,
-                                          0.5)
-
-    simulation.setScan(scan)
-    return simulation
-
+import PolarizedSpinAsymmetry as psa
 
 def run_simulation(q_axis, fitParams, *, polarizer_dir, analyzer_dir):
     """
@@ -96,115 +26,18 @@ def run_simulation(q_axis, fitParams, *, polarizer_dir, analyzer_dir):
     """
     parameters = dict(fitParams, **fixedParams)
 
-    sample = get_sample(parameters)
-    simulation = get_simulation(q_axis, parameters, polarizer_dir,
-                                analyzer_dir)
+    sample = psa.get_sample(parameters)
+    simulation = psa.get_simulation(q_axis, parameters, polarizer_dir,
+                                    analyzer_dir)
 
     simulation.setSample(sample)
     simulation.runSimulation()
     return simulation
 
-
-def qr(result):
-    """
-    Returns two arrays that hold the q-values as well as the
-    reflectivity from a given simulation result
-    """
-    q = numpy.array(result.result().convertedBinCenters(ba.Axes.QSPACE))
-    r = numpy.array(result.result().array(ba.Axes.QSPACE))
-
-    return q, r
-
-
-####################################################################
-#                         Plot Handling                            #
-####################################################################
-
-
-def plot(qs, rs, exps, labels, filename):
-    """
-    Plot the simulated result together with the experimental data
-    """
-    fig = plt.figure()
-    ax = fig.add_subplot(111)
-
-    for q, r, exp, l in zip(qs, rs, exps, labels):
-
-        ax.errorbar(exp[0],
-                    exp[1],
-                    xerr=exp[3],
-                    yerr=exp[2],
-                    fmt='.',
-                    markersize=0.75,
-                    linewidth=0.5)
-
-        ax.plot(q, r, label=l)
-
-    ax.set_yscale('log')
-    plt.legend()
-
-    plt.xlabel("Q [nm${}^{-1}$]")
-    plt.ylabel("R")
-
-    plt.tight_layout()
-    plt.savefig(filename)
-
-
-def plotSpinAsymmetry(data_pp, data_mm, q, r_pp, r_mm, filename):
-    """
-    Plot the simulated spin asymmetry as well its
-    experimental counterpart with errorbars
-    """
-
-    # compute the errorbars of the spin asymmetry
-    delta = numpy.sqrt(4 * (data_pp[1]**2 * data_mm[2]**2 + \
-            data_mm[1]**2 * data_pp[2]**2 ) /
-                ( data_pp[1] + data_mm[1] )**4 )
-
-    fig = plt.figure()
-    ax = fig.add_subplot(111)
-
-    ax.errorbar(data_pp[0],
-                (data_pp[1] - data_mm[1])/(data_pp[1] + data_mm[1]),
-                xerr=data_pp[3],
-                yerr=delta,
-                fmt='.',
-                markersize=0.75,
-                linewidth=0.5)
-
-    ax.plot(q, (r_pp - r_mm)/(r_pp + r_mm))
-
-    plt.gca().set_ylim((-0.3, 0.5))
-
-    plt.xlabel("Q [nm${}^{-1}$]")
-    plt.ylabel("Spin asymmetry")
-
-    plt.tight_layout()
-    plt.savefig(filename)
-
-
-####################################################################
-#                          Data Handling                           #
-####################################################################
-
-
-def load_exp(fname):
-    dat = numpy.loadtxt(fname)
-    return numpy.transpose(dat)
-
-
-def filterData(data, qmin, qmax):
-    minIndex = numpy.argmin(numpy.abs(data[0] - qmin))
-    maxIndex = numpy.argmin(numpy.abs(data[0] - qmax))
-
-    return data[:, minIndex:maxIndex + 1]
-
-
 ####################################################################
 #                          Fit Function                            #
 ####################################################################
 
-
 def run_fit_ba(q_axis, r_data, r_uncertainty, simulationFactory,
                startParams):
 
@@ -241,8 +74,8 @@ if __name__ == '__main__':
     datadir = os.getenv('BORNAGAIN_EXAMPLE_DATA_DIR', '')
     fname_stem = os.path.join(datadir, "MAFO_Saturated_")
 
-    expdata_pp = load_exp(fname_stem + "pp.tab")
-    expdata_mm = load_exp(fname_stem + "mm.tab")
+    expdata_pp = psa.load_exp(fname_stem + "pp.tab")
+    expdata_mm = psa.load_exp(fname_stem + "mm.tab")
 
     if len(argv) > 1 and argv[1] == "fit":
         fixedParams = {
@@ -291,17 +124,17 @@ if __name__ == '__main__':
                               polarizer_dir=ba.R3(0, -1, 0),
                               analyzer_dir=ba.R3(0, -1, 0))
 
-    qzs = numpy.linspace(qmin, qmax, scan_size)
-    q_pp, r_pp = qr(run_Simulation_pp(qzs, paramsInitial))
-    q_mm, r_mm = qr(run_Simulation_mm(qzs, paramsInitial))
+    qzs = numpy.linspace(psa.qmin, psa.qmax, psa.scan_size)
+    q_pp, r_pp = psa.qr(run_Simulation_pp(qzs, paramsInitial))
+    q_mm, r_mm = psa.qr(run_Simulation_mm(qzs, paramsInitial))
 
-    data_pp = filterData(expdata_pp, qmin, qmax)
-    data_mm = filterData(expdata_mm, qmin, qmax)
+    data_pp = psa.filterData(expdata_pp, psa.qmin, psa.qmax)
+    data_mm = psa.filterData(expdata_mm, psa.qmin, psa.qmax)
 
-    plot([q_pp, q_mm], [r_pp, r_mm], [data_pp, data_mm], ["$++$", "$--$"],
+    psa.plot([q_pp, q_mm], [r_pp, r_mm], [data_pp, data_mm], ["$++$", "$--$"],
          "MAFO_Saturated_initial.pdf")
 
-    plotSpinAsymmetry(data_pp, data_mm, qzs, r_pp, r_mm,
+    psa.plotSpinAsymmetry(data_pp, data_mm, qzs, r_pp, r_mm,
                       "MAFO_Saturated_spin_asymmetry_initial.pdf")
 
     if fit:
@@ -313,14 +146,14 @@ if __name__ == '__main__':
         print("Fit Result:")
         print(fitResult)
 
-        q_pp, r_pp = qr(run_Simulation_pp(qzs, fitResult))
-        q_mm, r_mm = qr(run_Simulation_mm(qzs, fitResult))
+        q_pp, r_pp = psa.qr(run_Simulation_pp(qzs, fitResult))
+        q_mm, r_mm = psa.qr(run_Simulation_mm(qzs, fitResult))
 
-        plot([q_pp, q_mm], [r_pp, r_mm], [data_pp, data_mm],
+        psa.plot([q_pp, q_mm], [r_pp, r_mm], [data_pp, data_mm],
              ["$++$", "$--$"], "MAFO_Saturated_fit.pdf")
         plt.draw()
 
-        plotSpinAsymmetry(data_pp, data_mm, qzs, r_pp, r_mm,
+        psa.plotSpinAsymmetry(data_pp, data_mm, qzs, r_pp, r_mm,
                           "MAFO_Saturated_spin_asymmetry_fit.pdf")
         plt.draw()