load'files'
NB. get_ions (boxed list) --> reads in atomic data
NB. the arrays are assigned to nouns with the listed names
get_ions=: monad define
n=. #y                  NB. number of ions in list
while. (n=. <:n)>: 0 do.  
(n{y)=: readXn n{y
end.
print 'Atomic data read.'
)

idx=: # i.@# 
gdn=: \:~               NB. grade down 
mark=: <'**********'    NB. separator of data blocks
upper_tri=: -.@>:/~@i.  NB. mask off lower triangular part of matrix

NB. Read in data for specific ion 'Xn':
readXn=: monad define
Xn=. >y                          NB. read data file and box each line:
zz=. < ;. 2 fread 'ATOMIC_DATA/',Xn,'.dat'
section=. idx mark = }: &. >zz   NB. get locations of separator marks
z1=. (i. 0{section) { zz                  NB. levels & their energies
z2=. ((>: 0{section)}. i. 1{section){ zz  NB. Einstein A's
z3=. ((>: 1{section)}. i. 2{section){ zz  NB. Collision cross-sections
   NB. the blank lines have a single element: < LF=. 10{a.  
nonblank =: idx -. 1 = ># &. > z1 
z1=. nonblank { z1                 NB. discard the blank (LF only) lines
lname=. >0{"1 ;: >z1               NB. extract the level designation string
numbs=. ".  >1{"1 ;: >z1           NB. extract the numeric part
size=. # EeV=. 1{"1 numbs          NB. second column is energy in eV
z2=. ". > 0{"1 ;: >z2              NB. get i,j, and A[i->j] values
ij=. ,<"1 <: 0 1{"1 z2             NB. extract (i,j), convert to boxed (i-1,j-i)
A=. |:(2{"1 z2) ij}(size,size)$0   NB. fill matrix with appropriate A's
wt=. |: (size,size)$ 0{"1 numbs    NB. matrix with statistical weights as rows
E=. (upper_tri size)*(|: EeV-/EeV) NB. make matrix of energy differences
('OM_',Xn)=: ". > 0{"1 ;: >z3      NB. store Omega table as e.g. OM_O3
('Lname_',Xn)=: lname              NB. store electron configs
   NB. fill grand matrix with As,Es & wts. Omegas (T dependent) will go in later.
(>A;E;wt) (1 2 3) }(4,size,size)$0 
)

ion_list=: 'C1';'C2';'C3';'C4';'N1';'N2';'N3';'N4';'N5'
ion_list=: ion_list,('O1';'O2';'O3';'O4';'O5';'O6')
ion_list=: ion_list,('Fe2';'Ne2';'Ne3';'Ne4';'Ne5')

Omegas=: dyad define       NB. Usage: 'Xn' Omegas T -> Omega(T) table
z3=. ". ('OM_',x)          NB. get the table OM_Xn
t=. _2}. 0{ z3             NB. 1st row is tabulated temps; discard trailing 0's
dat=. }. z3                NB. remove 1st row
ij=. ,<"1 <: 0 1{"1 dat    NB. 1st 2 columns are (i,j) indices; box them
dat=. |: 2}."1 dat         NB. trim off index columns and transpose array
size=. >: {. gdn ,>ij      NB. highest level (+ 1) = size of array
t0=. 1e_4 * y              NB. convert from T to t = 10e-4 * T
OM=. (t;dat) splint0 t0    NB. interpolate Omega for each (i->j)
|: OM ij} (size,size)$0    NB. form results into matrix
)

NB. Matrix of (collisional+radiative) rates for all i->j
rate=: dyad define     NB. Usage: (ion_name) rate (T N_e)
'Om A E w'=. x         NB. separate components of data matrix
'T Ne'=. y 
t=. 1e_4*T             NB. convert from T to t = 10e-4 * T
d=. 8.629e_8*Om% %:t   NB. downward collisional rates (/electon/ion)
u=. (w%~|:d)*(^_1.1605*E%t)  NB. upward collisional rates
A + Ne*u + d%w         NB. up coll & down coll + Einstein As
)

pd=: (<0 1)&|:   NB. get the principal diagonal of a matrix
unit=: =@i.      NB. make an identity matrix
norm=: ] % +/    NB. normalize a population vector to unity
apply=: [: norm +/ .*     NB. matrix times vector, then normalize
biavg=: 3 : '-: +/ >y'    NB. average of two boxed vectors
iterate=: ([: }. ]) , [: < [ apply [: biavg ]  

NB. Solve statistical equilibrium equations for level populations.
pop=: dyad define      NB. Usage: 'ion_name' pop (T,N_e)               
Xn=. ". x              NB. store data array whose name is 'x' in Xn
XOM=. x  Omegas 0{y    NB. interpolate in Omega table for Omega(T)
Xn=. XOM 0} Xn         NB. put Omega(T) into the grand Xn data array
n=. # M=. Xn rate y    NB. get the matrix of total transition rates
M=. (-|:M)+(=@i.n)* +/"1 M    NB. form statistical equilibrium eqns
M2=. -(-. unit n)* M%(pd M)   NB. construct a' matrix for iteration
   NB. Take the Boltzman distribution as the first guess:
p1=. pte=: norm (0{"1(3){Xn)*(^_11605*(0{2{Xn)%(0{y))    
if. (+/}. gdn pte)<1e_200 do.      NB. if excited level pop < 1e-200
bn=. pse =. pte=: 1 (0)} pte > 2   NB. zero out excited levels        
else.
p2=. M2 apply p1       NB. 1st iteration
bn=. pte %~ pse=. >1{ (M2&iterate)^:_ p1;p2  NB. iterate to convergence
end.
(3,n)$ pte,pse,bn  NB. Boltzmann, statistical equil., departure coeff b_n
)

emiss=: dyad define  NB. Matrix of radiative emission for all transitions
'Om A E w'=. ". x             NB. get data array
pse=. 1{ x pop y              NB. get level populations
1.6021765e_12*(|:E)*A*pse     NB. line emission = N_i * A_ij * E_ij
)

FAINT_LINE_LIMIT=: 2e_5  NB. limiting fractional intensity for line list
NB. reset this to zero to get *all* the lines.

NB. Compute emission line spectrum. Usage: 'ion_name' lines (T, N_e)
NB. output: (upper level)->(lower level), wavelength, emission (/ion)
lines=: dyad define
'Om A E w'=. ". x             NB. get data array for ion x
lname=. <"1 ". 'Lname_',x     NB. get names of atomic levels
a=. ,aa =. 0~:A               NB. boolean array: is Einstein A non-zero?
nrow=. |:ncol=. (i.#aa) (i.#aa)}aa  NB. matrices filled with row/col indices
e=. a * ,|:E    NB. e is E of transitions with non-0 A; \:e sorted down
i=.(+/a){. \:e  NB. (+/a) is is # of transitions with non-zero A-values.
z1=.a%1+0.000287*(1+567000%(a=.12398.42*%i{e)^2)  NB. Air wavelength correction.
z1=.(Q*a)+(-.Q=.a<2000)*z1    NB. Vacuum wavelengths below 2000A.
z1=. z1*(z1<10000)+ 1e_4*(z1>:10000) NB. lambda in microns if > 10,000 A. 
tz=. +/ z2=. i{ ,x emiss y    NB. get the line emission
bright=. idx z2 > FAINT_LINE_LIMIT*tz  NB. indices of lines brighter than limit
upr=. (i{,nrow){ lname
lwr=. (i{,ncol){ lname
arrow=. (#z2)#<'-->'
bright { |: (5,(#z2))$upr,arrow,lwr,(<"0 z1),<"0 z2
)

O3r=: 3 : '%/ (3 2;4 3){ ''O3'' emiss y'    NB. yeilds 5007%4363 ratio
O2r=: 3 : '%/ (1 0;2 0){ ''O2'' emiss y'    NB. yeilds 3729%3726 ratio

NB. Total radiative loss (/ion/electron). Usage: 'ion' Rad_Loss (T,N_e)
Rad_Loss=: 4 : '(1{y)%~ +/+/ x emiss y' 

get_ions ion_list

NB.================ spline interpolation code ==========================

NB. Spline interpolation: Usage (x;y) splint0 x0
NB. x0 may be an array
splint0=: dyad define
'xx yy'=. x
D=. xx spline yy
(xx;yy;D) spltrp0 y
)

NB. Find 2nd derivatives for spline fit.
NB. Usage: x spline y --> D
spline=: dyad define
h2=. }. h=. DEL x
a=. 0.5 - c=. -: h2 * hh=. % h2 + }:h
b=. 1#~#c     
u=. 3*hh*DEL (DEL y)%h
0,(tridgl a;b;c;u),0
)

NB. Usage: (xx;yy;d) spltrp0 x0 --> y(x0)
NB. xx, yy & d boxed; x0 may be an array
NB. Don't extrapolate; use last value if beyond ends.
spltrp0=: dyad define
'xx yy d'=. x
f0=. (}: xx)%(h=. DEL xx)
f1=. (}. xx)%h
g0=. (}: yy)%h
g1=. (}. yy)%h
g=. h*(g0*f1)-g1*f0
ff=. _1+ +/"1 y >/xx
t=. ff = n=. (0>. ff)<._2+#xx NB. t=0 beyond ends where n not = ff
lin=. (n{g)+ y*n{(g1-g0)      NB. linear interpolation
xn=. y % hn=. n{h
a=. (n{f1)-xn
b=. xn-(n{f0)
aa=. a* <: *:a                NB. a*(a^2 - 1) = a^3 - a
bb=. b* <: *:b
cub=. 6%~(*:hn)*(aa*n{}:d)+bb*n{}.d    NB. cubic term
nn=. n + (n<ff)               NB. use upper end of top interval
((1-t)*(nn{yy))+t*(lin+cub)
)

NB. Solve tridiagonal system by direct substitution.
NB. Rows are  a_i*x_(i-1)+b_i*x_i+c_i*x_(i+1) = u_i 
NB. a,b,c,u have same length, i.e., a[0]=0 and c[n]=0. 
NB. Usage:  soln <-- tridgl a;b;c;u
NB. 
tridgl=: monad define
'a b c u'=. y
n1=. <: n=. <: #b
x=. %{.b
c=. (x*{.c) 0}c
u=. (x*{.u) 0}u
j=. 0
while. n1>:j=.>:j do. 
  b1=. j{b
  a1=. j{a
  c0=. (<:j){c
  x=. %b1-a1*c0
  c1=. j{c
  c=. (x*c1)j}c
  u1=. j{u
  u0=. (<:j){u
  u=. (x*u1-a1*u0)j}u
end.
u=. (((n{u)-(n{a)*(n1{u))%(n{b)-(n{a)*(n1{c))n}u
j=. n
while. 0<:j=.<:j do.
  u1=. j{u
  c1=. j{c
  u2=. (>:j){u
  u=. (u1-c1*u2)j}u
end.
)