generate_datasim.py

# 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")

# Remove weirdcoeff and set to nm/DAC
#respmat= np.round(respmat/1.77/1717.412/170*10*0x7FFF*1e-6*2**30).astype('i4')

# 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=14

#############################################################################
# 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(-1e5, 1e5, 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:]
C_N_RND = 14
C_N_FP = 16
SUMSAT=36


K_A = 4950
K_B = -4000
K_C = 16384
K_D = 16300


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

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

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

# Step 3: correction computation
corrfp = np.zeros((N_PSC, size), dtype="i8")


corrfp[:,0] = mm[:,0] *K_A*2**-C_N_RND *K_C*2**-C_N_RND

for i in range(1,size):
    corrfp[:,i] = (mm[:,i]       *K_A*2**-C_N_RND +
                   mm[:,i-1]     *K_B*2**-C_N_RND +
                   corrfp[:,i-1] *K_D*2**-C_N_RND
                   )*K_C*2**-C_N_RND
                    
corr = np.round(corrfp*2**-C_N_FP).astype('i2')

#############################################################################
# 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("corrout.txt", "w") as fp:
    for i in range(size):
        fp.write(" ".join(corr[:,i].astype("str"))+"\n")