clc;
% This script provides a guide on how to extract designs from the dataset 
% and reshape them into spatial density fields for machine learning 
% applications.

file = fullfile('data.h5');

%% Specify sample index 
idx = 1;

%% Extract sample parameters
nx = double(h5read(file,['/',num2str(idx),'/nx']));         % number of elements in x-direction
ny = double(h5read(file,['/',num2str(idx),'/ny']));         % number of elements in y-direction
nElem = h5read(file,['/',num2str(idx),'/nElem']);           % total number of elements
nNodes = h5read(file,['/',num2str(idx),'/nNodes']);         % total number of nodes 
scale = h5read(file,['/',num2str(idx),'/scale']);           % Scaling parameters for the L-beam design domain

%% Extract designs
x = double(h5read(file,['/',num2str(idx),'/design/x']));    % non-fail-safe design
y = double(h5read(file,['/',num2str(idx),'/design/y']));    % fail-safe design

%% Determine element coordinates 
elemCoords = elementCoordinates(nElem,nx,ny,scale);         % element coordinates [x,y,z]

%% Reshape design from vector to matrix 
x = densityVec2Field(x,elemCoords,nx,ny);                   % non-fail-safe design            
y = densityVec2Field(y,elemCoords,nx,ny);                   % fail-safe design

%% Plot designs 
subplot(2,1,1)
colormap(gray); imagesc(1-rot90(x)); axis equal; axis off; title('non-fail-safe design'); drawnow; 
subplot(2,1,2)
colormap(gray); imagesc(1-rot90(y)); axis equal; axis off; title('fail-safe design'); drawnow; 

%% Functions
function elemCoords = elementCoordinates(nElem,nx,ny,scale)
    if nElem < nx*ny
        % first rectangle
        X1 = repmat(1:nx*scale(1),1,ny);
        Y1 = repmat(1:ny,nx*scale(1),1);
        Z1 = zeros(nx*scale(1)*ny,1);
        
        % second rectangle
        X2 = repmat(nx*scale(1)+1:nx,1,ny*scale(2));
        Y2 = repmat(1:ny*scale(2),nx-nx*scale(1),1);
        Z2 = zeros((nx-nx*scale(1))*ny*scale(2),1);
        
        % concatenate
        X = [X1';X2'];
        Y = [Y1(:);Y2(:)];
        Z = [Z1;Z2];
        elemCoords = [X,Y,Z];
    else 
        X = repmat(1:nx,1,ny);
        Y = repmat(1:ny,nx,1);
        Z = zeros(nx*ny,1);
        elemCoords = double([X',Y(:),Z]);
    end
end

function vField = densityVec2Field(v,elemCoords,nx,ny) 
    if  numel(v) < nx*ny
        vField = zeros(nx,ny);
        for i = 1:size(elemCoords,1)
            vField(index(i,1),index(i,2)) = v(i);
        end
    else
        vField = reshape(v,nx,ny);
    end
end