%% Exercise Session 2:  Plasma current and position reconstruction/determination
%% Load the data and get plasma grid
% As a first step we need to get the data (geometries, matrices, measurements 
% etc.)

clear;
fromfile = true; % set true to load from file
fname = 'ex2_data.mat'; 
if fromfile && ~isempty(which(fname))
  load('ex2_data.mat'); % if file exists, load it
else
  % Loading an equilibrium from TCV
  [L,~,LY] = fbt('tcv',61400,1,'selu','v','izgrid',1);
  % uncomment this to save file
  save(fname,'L','LY');
end
%% Define plasma distribution on the x grid

Ix = zeros(L.nzx,L.nrx); Ix(2:end-1,2:end-1) = LY.Iy;
% true current centroid
r0 = sum(L.rrx(:).*Ix(:))/sum(Ix(:));
z0 = sum(L.zzx(:).*Ix(:))/sum(Ix(:));
Ip0=sum(Ix(:)); %calculating the total current 
% plot distribution and measurement locations

figure(1); clf;
imagesc(L.rx,L.zx,Ix); axis tight; axis xy equal; %plots current distribution
hold on;
plot(L.G.rl,L.G.zl);
plot(L.G.rv,L.G.zv); 
plot(r0,z0,'wo');% current centroid
colorbar;
legend('limiter','vessel','');
title('Example of a current distribution');

%% Compute expected measurements for given plasma distribution
Ia = LY.Ia; % circuit currents
% measurements = plasma current contribution + circuit current contribution
Bm  = L.G.Bmx*Ix(:) + L.G.Bma*Ia; % magnetic probes 
Ff  = L.G.Mfx*Ix(:) + L.G.Mfa*Ia; % flux loops

%% Add noise and turn into time-dependent signal
% In reality, it is a hard to clear measurements during operation. Usually, 
% the measurements also contain noise. Hence, we add time-dependent noise to the 
% magnetic fields and flux measurements.

eBm = 10e-3; % 10mT noise
eFf = 10e-3; % 10mWb noise
eIa = 0;     % assume no measurement error on currents

% time grid
time = 0:1e-3:0.1;
nt = numel(time); 
% time evolution
Bm_meas = Bm + eBm*randn(numel(Bm),nt);
Ff_meas = Ff + eFf*randn(numel(Ff),nt);
Ia_meas = Ia + eIa*randn(numel(Ia),nt);
%% Plots

figure(2); clf;
% magnetic probes time evolution
subplot(221)
plot(time,Bm_meas);
title('B probe [T]'); xlabel('time [s]')
% flux loops time evolution
subplot(223)
plot(time,Ff_meas);
title('\psi_{fluxloop} [Wb]'); xlabel('time [s]')

% measurements per magnetic probe
subplot(222)
plot(1:numel(Bm),Bm_meas,'.',1:numel(Bm),Bm,'or')
title('B probes [T]')
xlabel('probe number');

% measurements per flux loop
subplot(224)
plot(1:numel(Ff),Ff_meas,'.',1:numel(Ff),Ff,'or');
title('flux loops [Wb]')
xlabel('flux loop');

%Positons of the magentic probes and flux loops
figure(3); clf;
subplot(121)
imagesc(L.rx,L.zx,Ix);  axis tight; axis xy equal; colorbar ('westoutside');
hold on;
plot(L.G.rl,L.G.zl);
plot(L.G.rv,L.G.zv); 
plot(L.G.rf,L.G.zf,'kx');
text(L.G.rf+0.02,L.G.zf,L.G.dimf);
legend('limiter','vessel','Flux loop');
title('flux loop in TCV');

subplot(122)
imagesc(L.rx,L.zx,Ix); axis tight; axis xy equal; colorbar('westoutside');
hold on;
plot(L.G.rl,L.G.zl);
plot(L.G.rv,L.G.zv); 
plot(L.G.rm,L.G.zm,'kx')
text(L.G.rm+0.02,L.G.zm,L.G.dimm);
legend('limiter','vessel','B probe');
title('magnetic probes in TCV');

%% Define new grid boundaries of representation with fewer elements (2x3)
% In the following, we demonstrate how define a current distribution with 
% fewer degrees of freedom. We split the |x| grid into several rectangular
% regions and let the current distribution be constant in that region.

nr = 2; nz = 3; nh = nr*nz; %number of regions in radial and vertical direction; total number of degrees of freedom
rhgrid = linspace(min(L.G.rx),max(L.G.rx),nr+1);
zhgrid = linspace(min(L.G.zx),max(L.G.zx),nz+1);

%% Build transformation matrix $$\mathbf{T_{xh}}$$ from full 'x' grid and a reduced 'h' representation
% In this section, we compute a matrix $\mathbf{T_{xh}}$, such that  ${I_x}  
% =\mathbf{T_{xh}}\cdot{I_h}$. Each column of |Txh| defines a
% current distribution that is piecewise constant in a region of the |x|
% grid and zero otherwise.

Txh = zeros(L.nx,nh);
[rrx,zzx] = meshgrid(L.G.rx,L.G.zx);
for ir = 1:nr
    for iz = 1:nz
        ii = (ir-1)*(nz) + iz; % element index
        % choose region bounded by points defining borders of the rectangles
        txy = (rhgrid(ir) <= rrx) & (rrx < rhgrid(ir+1))...
            & (zhgrid(iz) <= zzx) & (zzx < zhgrid(iz+1)) ;
        Txh(:,ii) = txy(:);
    end
end
% calculate mean r,z of each new grid point for illustration
rhmean = sum(Txh.*rrx(:),1)./sum(Txh,1);
zhmean = sum(Txh.*zzx(:),1)./sum(Txh,1);

% plot 
figure(4); clf;
subplot(121)
p=plot(rrx,zzx,'.b','markersize',2); hold on; plot(rhmean,zhmean,'or'); axis equal tight
[rhh,zhh] = meshgrid(rhgrid,zhgrid);
hold on; p1=plot(rhh,zhh,'+k');
legend([p(1) p(29) p1(1)],'x grid (filaments)','h grid midpoints','h grid boundary','location','south')
title('x and h grid definitions')
% example of a distribution at the smaller h-grid
Ih = zeros(nh,0); Ih([1,3,6],1) = 1e3;
Ix = Txh*Ih;
subplot(122)
pcolor(L.G.rx,L.G.zx,reshape(Ix,L.nzx,L.nrx));colorbar; 
title('Example of an (unlikely) current distribution on the h grid')
%% Pretask
% Before starting the exercise, run the code and get familiar with output.


%% a) Least-squares problem for current distribution on the full grid 
% Write a least-squares problem to determine the current distribution $I_x$ 
% on the full $\mathbf{x}$ grid from the measurements of magnetic probes and 
% 
% flux loops, in the form $\min||AI_x - b||_2$. Can this 
% least-squares problem be solved? If not, why not?

%Enter your solution here


%% 
% 
%% b) Least-squares problem for current distribution at reduced h-grid
% Write a least-squares problem to determine the current distribution $I_h$ 
% on the reduced grid from the measurements. What is $A$ and $b$? Can this problem be solved? If not, 
% why not?

%Enter your solution here


%% 
% 
%% c) Solution of the least-squares problem
% Write the solution of least-squares problem in (b) in the form $I_h = Hb$. 
% How can be $H$ determined from $A$?

%Enter your solution here


%% 
% 
%% d) Write a linear estimator for the plasma current 
% Recall that the plasma current can be determined from the current distribution 
% as $I_p=\sum_{j=1}^{n_x}I_{x,j}$. Write a linear estimator for the plasma current 
% from the measurements such that $I_p=H_{I_p}\cdot b$.

%Enter your solution here


%% 
% 
%% e) Estimate of r and z based on the measurements
% We can define an estimator for the plasma radial and vertical position, by 
% assuming the position corresponds to the centroid of the current distribution:
% 
% $\hat{r}=\frac{\sum_{j}^{n_x}r_{x,j}I_{x,j}}{\sum_{i}^{n_x}I_{x,i}}$ and $\hat{z}=\frac{\sum_{j}^{n_x}z_{x,j}I_{x,j}}{\sum_{i}^{n_x}I_{x,i}}$
% 
% Use these experessions to compute an estimate of $r$ and $z$ based on the 
% measurements of magnetic probes and flux loops. Is this a linear estimator?

%Enter your solution here


%% 
% 
%% f) Estimator for the plasma $rI_p$ and $zI_p$ 
% Write a estimator for the plasma $rI_p$ and $zI_p$ i.e. the product of radial/vertical 
% position and plasma current. Is this linear?

%Enter your solution here



%% g) Outputs of the linear estimators
% Plot the outputs of the linear estimators and compare to the known values 
% of $r_0I_{p0},z_0I_{p0},I_{p0}$. Improve the estimate of $z_0I_{p0}$ and  $I_{p0}$ and explain 
% your solution. 

%Enter your solution here

%% 
% 
%% h) Effect of nh
% Plot the mean and standard deviation of $I_p$, $zI_p$, $rI_p$ estimates for % various values of $nr,nz$ (and $nh=nz*nr$).
% Does the estimator quality increase, decrease or stay the same when increasing $nh$? Explain why.

%Enter your solution here


%% 
% 
%% i) Effect on Circuit current measurements
% Now assume that the circuit current measurements \texttt{Ia\_meas} is also 
% affected by measurement noise of standard deviation \texttt{eIa=100[A]}. 
% Check the quality of your estimator and compare it to the non-noisy case. 
% Reformulate the least squares problem to jointly estimate the plasma current distribution 
% and circuit currents from the measurements and compare the estimation results.
% 
% 
% 
% 
% 
%