1   # stations list
   2   STATIONS = range(1,15);
   3   # 3 sources
   4   SOURCES = ('a','b','c');
   5 
   6   #--- these are generated automatically from the station list
   7   # list of ifrs as station pairs: each entry is two indices, (s1,s2)
   8   IFRS   = [ (s1,s2) for s1 in STATIONS for s2 in STATIONS if s1<s2 ];
   9   # list of ifr strings of the form "s1-s2"
  10   QIFRS  = [ '-'.join(map(str,ifr)) for ifr in IFRS ];
  11   # combined list: each entry is ("s1-s2",s1,s2)
  12   QQIFRS = [ ('-'.join(map(str,ifr)),) + ifr for ifr in IFRS ];
  13 
  14   # namespace of all node classes
  15   Meq = Namespace('Meq');
  16   # global node scope & repository
  17   nsglob = NodeScope();
  18   # init some node groups
  19   ROOT = NodeGroup('root');
  20   SOLVERS = NodeGroup('solvers');
  21 
  22   #------- create nodes for instrumental model
  23   # some handy aliases
  24   ZERO = nsglob.zero() << Meq.Constant(value=0);
  25   UNITY = nsglob.unity() << Meq.Constant(value=1);
  26 
  27   PHASE_CENTER_RA  = nsglob.ra0() << Meq.Parm();
  28   PHASE_CENTER_DEC = nsglob.dec0() << Meq.Parm();
  29 
  30   STATION_UWV = {};
  31   STATION_POS = {};
  32 
  33   ARRAY_POS = nsglob.xyz0() << Meq.Composer(
  34     nsglob.x0() << Meq.Parm(),
  35     nsglob.y0() << Meq.Parm(),
  36     nsglob.z0() << Meq.Parm() );
  37 
  38   # create station-related nodes and branches
  39   for s in STATIONS:
  40     STATION_POS[s] = nsglob.xyz(s) << Meq.Composer(
  41       nsglob.x(s) << Meq.Parm(),
  42       nsglob.y(s) << Meq.Parm(),
  43       nsglob.z(s) << Meq.Parm() );
  44     STATION_UWV[s] = nsglob.stuwv(s) << Meq.UWV(children={
  45       'xyz': STATION_POS[s],
  46       'xyz0': ARRAY_POS,
  47       'ra': PHASE_CENTER_RA,
  48       'dec': PHASE_CENTER_DEC
  49     });
  50     # create per-source station gains
  51     for (q,src) in enumerate(SOURCES):
  52       nsglob.G(s,q=q) << Meq.Composer(
  53         nsglob.Gxx(s,q=q) << Meq.Polar(Meq.Parm(),Meq.Parm()), ZERO,
  54         ZERO, nsglob.Gyy(s,q=q) << Meq.Polar(Meq.Parm(),Meq.Parm()),
  55       dims=[2,2]);
  56     # alternative: single gain with no direction dependence
  57     # nsglob.G(s) << Meq.Composer(
  58     #    nsglob.Gxx(s) << Meq.Polar(Meq.Parm(),Meq.Parm()), ZERO,
  59     #    ZERO, nsglob.Gyy(s) << Meq.Polar(Meq.Parm(),Meq.Parm()),
  60     #  dims=[2,2]);
  61 
  62   # this function returns a per-station gain node, given a set of qualifiers
  63   def STATION_GAIN (s=s,q=q,**qual):
  64     # **qual swallows any remaining qualifiers
  65     return nsglob.G(s,q=q);
  66     # note alternative for no direction dependence:
  67     # def STATION_GAIN (s=s,**qual):  
  68     #   return nsglob.G(s);
  69 
  70   #------- end of instrumental model
  71 
  72   #------- create model for unpolarized point source
  73   # References instrumental model: STATION_GAIN(s,**qual), STATION_UWV[s].
  74   # Returns unqualified predict node, should be qualified with ifr string.
  75   def makeUnpolarizedPointSource (ns,**qual):
  76     ns.lmn() << Meq.LMN(children={
  77         'ra':   ns.ra() << Meq.Parm(),
  78         'dec':  ns.dec() << Meq.Parm(),
  79         'ra0':  PHASE_CENTER_RA,
  80         'dec0': PHASE_CENTER_DEC
  81     });
  82     ns.stokes_i() << Meq.Parm();
  83     # create per-station term subtrees
  84     for s in STATIONS:
  85       ns.sdft(s) << Meq.MatrixMultiply(
  86         STATION_GAIN(s,**qual),
  87         Meq.StatPSDFT(ns.lmn(),STATION_UWV[s])
  88       );
  89     # create per-baseline predicters
  90     for (q,s1,s2) in QQIFRS:
  91       ns.predict(q) << Meq.Multiply(
  92           ns.stokes_i(),
  93           ns.dft(q) << Meq.DFT(ns.sdft(s1),ns.sdft(s2)),
  94       );
  95     return ns.predict;
  96 
  97   #------- create peeling unit
  98   # inputs: an unqualified input node, will be qualified with ifr string.
  99   # predicters: list of unqualified predict nodes, will be qualified with ifr string.
 100   # Returns unqualified output node, should be qualified with ifr string.
 101   def peelUnit (inputs,predicters,ns):
 102     for q in QIFRS:
 103       # create condeq branch
 104       ns.condeq(q) << Meq.Condeq(
 105         ns.measured(q) << Meq.PhaseShift(children=inputs(q)),
 106         ns.predicted(q) << Meq.Add(*[prd(q) for prd in predicters])
 107       );
 108       # create subtract branch
 109       ns.subtract(q) << Meq.Subtract(ns.measured(q),ns.predicted(q));
 110     # creates solver and sequencers
 111     ns.solver() << Meq.Solver(*[ns.condeq(q) for q in QIFRS]);
 112     for q in QIFRS:
 113       ns.reqseq(q) << Meq.ReqSeq(ns.solver(),ns.subtract(q));
 114     # returns root nodes of unit
 115     return ns.reqseq;
 116 
 117   # create source predictors, each in its own subscope
 118   predicter = {};
 119   for (q,src) in enumerate(SOURCES):
 120     predicter[q] = makeUnpolarizedPointSource(nsglob.subscope('predict',src),q=q);
 121 
 122   # create spigots
 123   for q in QIFRS:
 124     nsglob.spigot(q) << Meq.Spigot();
 125 
 126   # chain peel units, by connecting outputs to inputs. First input
 127   # is spigot.
 128   inputs = nsglob.spigot;
 129   for (q,src) in enumerate(SOURCES):
 130     ns_pu = nsglob.subscope('peelunit',q);
 131     inputs = peelUnit(inputs,predicter.values(),ns=ns_pu);
 132     SOLVERS << ns_pu.solver();
 133 
 134   # create sinks, connect them to output of last peel unit
 135   for q in QIFRS:
 136     ROOT << nsglob.sink(q) << Meq.Sink(inputs(q));
 137 
 138   # create data collectors (this simply shows off the use of arbitrary node
 139   # groupings)
 140   ROOT << nsglob.solver_collect() << Meq.DataCollect(*SOLVERS.values());
 141 
 142   # deliberately create an orphan branch. This checkes orphan collection.
 143   # this whole branch should go away, and so should the UNITY node, which
 144   # is not used anywhere
 145   nsglob.orphan() << Meq.Add(Meq.Const(value=0),UNITY,Meq.Add(UNITY,ZERO));
 146 
 147   # resolve all nodes
 148   nsglob.resolve(ROOT);