Skip to content
Snippets Groups Projects
generate_datasim.py 3.79 KiB
Newer Older
# This script generate simulation data for the PI corrector + matrix multiplication.
import numpy as np


#############################################################################
# Response matrix, ref orbit and data input model
#############################################################################

# The inverted response matrix, we use an actual true dataset
respmat = np.load("respmat.npy")
# write to file
with open("respmat.txt", "w") as fp:
    for r in respmat.flatten():
        fp.write("{}\n".format(r))

# The orbit reference, again this is from true data
refmat = np.load("reforbit.npy")
# write to file
with open("reforbit.txt", "w") as fp:
    for r in refmat:
        fp.write("{}\n".format(r))

# BPM data generation
def generate_bpm_data(size, DC=0, A=0, F=100, N=0):
    """ Generate dummy BPM data, based on a simple model of DC + A.sin(F) + N.noise """

    sig = DC + A*np.sin(2*np.pi*F*np.arange(size)*1e-4+np.random.rand()*2*np.pi) + np.random.default_rng().normal(0, N, size)

    return sig.astype("int32")

#############################################################################
# Parameters
#############################################################################
N_BPM=respmat.shape[1]*2
N_PSC=respmat.shape[0]
size=200

N_MM_RND=23

#############################################################################
# Generate data input
#############################################################################
bpmdata=np.empty((N_BPM, size), dtype="int32")
for i, DC, A, F, N in zip(
        np.arange(N_BPM, dtype="int"),
        np.random.default_rng().uniform(-1e6, 1e6, N_BPM),
        np.random.default_rng().uniform(0, 200e3, N_BPM),
        np.random.default_rng().uniform(2,2000, N_BPM),
        np.random.default_rng().uniform(0, 20e3, N_BPM),):
    bpmdata[i] = generate_bpm_data(size, DC, A, F, N)

# write to file
# Warning, write x then y for each bpm, not all x then all y as the array is packed.
with open("bpmdata.txt", "w") as fp:
    for i in range(size):
        for j in range(N_BPM//2):
            fp.write("{} {} ".format(bpmdata[j,i], bpmdata[j+N_BPM//2,i]))
        fp.write("\n")    
        #fp.write(" ".join(bpmdata[:,i].astype("str"))+"\n")

#############################################################################
# Compute correction
#############################################################################
# change respmat shape, tiled
trespmat = np.zeros((N_PSC, N_BPM), dtype="int64")
trespmat[:50,:122] = respmat[:50]
trespmat[50:,122:] = respmat[50:]

KP=2**16
KI=0

## -----------------------
# Model computation

# Step 1: error computation
oe=bpmdata-refmat[:,np.newaxis]

# Step 2: matrix multiplication
mm=np.empty((N_PSC, size), dtype="int32")
for i in range(size):
    # Stage 1 : matmul
    _mm= np.matmul(oe[:,i].astype("int64"), trespmat.T, dtype='int64')
    mm[:,i] = _mm>>N_MM_RND

# Step 3: correction computation
mp=np.empty((N_PSC, size), dtype="int64")
mi=np.empty((N_PSC, size), dtype="int64")
mc=np.empty((N_PSC, size), dtype="int16")

mp = KP*mm.astype("int64")
mi = (KI*np.cumsum(mm.astype("int64"), axis=1))>>8
mc=np.clip(np.round((mp+mi)/2**34), -2**16, 2**16-1).astype("int16")

#############################################################################
# Write partial and results to files
#############################################################################
with open("orbiterror.txt", "w") as fp:
    for i in range(size):
        fp.write(" ".join(oe[:,i].astype("str"))+"\n")

with open("matmult.txt", "w") as fp:
    for i in range(size):
        fp.write(" ".join(mm[:,i].astype("str"))+"\n")

with open("propcorr.txt", "w") as fp:
    for i in range(size):
        fp.write(" ".join(mp[:,i].astype("str"))+"\n")

with open("corrout.txt", "w") as fp:
    for i in range(size):
        fp.write(" ".join(mc[:,i].astype("str"))+"\n")