#  Python 3
#  Author: Jonathan Goodman, goodman@cims.nyu.edu
#  Fall 2024
#  for the class Scientific Computing
#  https://math.nyu.edu/~goodman/teaching/ScientificComputing2024/ScientificComputing.html

#  Create a fake data time series that is a sum of oscillations with ...
#  ... noise added

# This file: FakeData.py,

import numpy as np
import array                  # to write bindary output files

rng = np.random.default_rng(seed=17)

T = 10.    # length of "observation" interval for all runs

#  First run, one frequency

omega = np.array([  .8 ])
a     = np.array([ 1.  ])
b     = np.array([  .2 ])
n     = 50
sig   = .1
run1 = {"n":n, "sig":sig, "omega":omega, "a":a, "b":b, "name":"run1"}

#  Second run, easy

omega = np.array([ 2.8 ,  1.2 ])
a     = np.array([ 1.  ,  0.  ])
b     = np.array([ 0.  ,  2.  ])
n     = 100
sig   = .6
run2 = {"n":n, "sig":sig, "omega":omega, "a":a, "b":b, "name":"run2"}

#  Third run, harder, four oscillators

omega = np.array([ 5.7 , 2.8 , 1.2 ,  .8])   # each column represents ...
a     = np.array([ .5  ,  .1 , 1.  , 0. ])   # one of the oscillators
b     = np.array([ .2  , 3.  , 0.  , 2. ])
n     = 1000
sig   = 1.4
run3 = {"n":n, "sig":sig, "omega":omega, "a":a, "b":b, "name":"run3"}

runs = [run1, run2, run3]

for run in runs:

   n    = run["n"]
   sig  = run["sig"]
   om   = run["omega"]
   a    = run["a"]
   b    = run["b"]
   name = run["name"]
   nf   = len(om)     # number of frequencies
   dt   = T/(n-1)     # approx time between observations
   
   times = np.zeros([n])
   xvals = np.zeros([n])
   
   for k in range(n):
      t = k*dt
      x = sig*rng.normal()
      for i in range(nf):
         x = x + a[i]*np.cos(om[i]*t) + b[i]*np.sin(om[i]*t)
      times[k] = t
      xvals[k] = x
      
   InfoFileName = "info_"+name+".py"
   with open(InfoFileName, 'w', encoding="utf-8") as info:
      infoLine = "n   = {n:5d}\n".format(n=n)
      info.write(infoLine)
      
      infoLine = "sig = {sig:12.6f}\n".format(sig=sig)
      info.write(infoLine)
      
      infoLine = "nf  = {nf:5d}\n".format(nf=nf)
      info.write(infoLine)
      
   DataFileName = "data_"+name               # for example, data_run1
   with open(DataFileName, 'wb') as DF:      # 'wb' is for "write binary"

      cpT = array.array('d', times)  # put the data in core Python arrays
      cpx = array.array('d', xvals)
      cpT.tofile(DF)                 # write it in binary to the data file
      cpx.tofile(DF)