# Improving accuracy of Neural Network for determining qubit rotation angle using Mathematica

The physics example problem (to illustrate the use of a basic neural network using Mathematica) I am looking at is a qubit rotated about the y-axis, where the rotation angle is discretized as $$\theta_j \in (0, \pi)$$. The setup involves the y-rotated qubit measured in the z-basis (hence spin-up and spin-down projector measurements). This scenario involves first analytically determining the measurement outcome probabilities as a function of the rotation angle $$\theta$$, then generating measurement outcomes for training, for specific fixed rotation angles $$\theta_j$$. Then generating another set of test measurement data for some fixed rotation angle $$\theta$$. We can then use the neural network to infer the most probable rotation angle. My training data involves generating m = 1000 total measurements for each discrete rotation angle $$\theta_j \in [0, \pi]$$, then saving the measurement outcomes as tuples of spin-up and spin-down outcomes for each discrete angle. These outcomes are associated with each of the discrete $$\theta_j$$ values which are one-hot vectors (hence training data of the form {1000,0} -> {1,0,0,0,0...} if for the first rotation angle we get all spin-up outcomes).

The idea is that after training, setting some true rotation angle $$\theta$$, and generating a new set of test measurement outcomes, the trained neural network should be able to output a probability distribution that shows the most likely discrete rotation angle is the true angle. The code below works but I am having difficulty improving the accuracy without simply increasing the layers and MaxTrainingRounds (this seems to have it's limits in improving accuracy). Can anyone advise on how to improve the accuracy of the code in determining the correct discrete rotation angle (I would like to maintain the general framework of the code)? I am very new to using Mathematica for machine learning applications hence the query. Thanks for any assistance, this is the code in question:

(*Spin-up/spin-down measurement probabilities on Y-rotated qubit *)
Pu[\[Theta]_] := (Cos[\[Theta]/2])^2;
Pd[\[Theta]_] := 1 - Pu[\[Theta]]

(*Discretize rotation values*)
theta = Table[\[Theta], {\[Theta], \[Pi]/100, \[Pi], \[Pi]/100}]

(*One-hot vectors from discretized theta values*)
numBins = Length[theta];
binToOneHot[bins_List, totalBins_Integer] :=
Table[If[i == #, 1, 0], {i, 1, totalBins}] & /@ bins
bins = Range[numBins];
oneHotEncoded = binToOneHot[bins, Length[theta]];

(*Evaluate measurement probabilities for each theta value*)
pr = Table[
List[N[FullSimplify[Re[Pu[\[Theta]]]]],
N[FullSimplify[Re[Pd[\[Theta]]]]]], {\[Theta], \[Pi]/
100, \[Pi], \[Pi]/100}]

(*Generate training measurement outcomes for each Subscript[\[Theta], \
j] (with random number generator) - 1 represents spin-up outcome and \
0 spin-down *)
(*Number of measurements for each Subscript[\[Theta], j] in (0,\[Pi])*)
m = 1000;
meas = Table[{}, {Length[pr]}];
Table[Do[ b = ConstantArray[0, 1]; Clear[r];
updateTally[] := Block[{r}, r = RandomReal[];
If[r < pr[[j]][[1]], b[[1]] += 1,
If[pr[[j]][[1]] <= r < pr[[j]][[1]] + pr[[j]][[2]],
b[[1]] += 0]]];
Do[updateTally[], {1}]; AppendTo[meas[[j]], b], {i, 1, m}], {j,
Length[pr]}];

(*Generated training measurement outcomes for each \
Subscript[\[Theta], j] *)
meas = Partition[Flatten[meas], m];

(*Function to count 1s and 0s*)
countOnesZeros[list_] := {Count[list, 1], Count[list, 0]}

(*Training data for neural network*)
(*Simulated Measurements Outcomes for each Subscript[\[Theta], j] in \
(0,\[Pi])*)

meas = Map[countOnesZeros, meas]
trainingData = Rule @@@ Transpose[{meas, oneHotEncoded}];

(*Neural network architecture*)
net = NetChain[{LinearLayer[50], ElementwiseLayer["ReLU"],
LinearLayer[100], SoftmaxLayer[]}];

(*Training the network*)
trainedNet =
NetTrain[net, trainingData, MaxTrainingRounds -> 50000,
LossFunction -> CrossEntropyLossLayer["Probabilities"]]

(*Example set of test measurements*)
meas2 = {meas[[12]]}
trainedNet /@ meas2

(*Probability unity check*)
Total[(trainedNet /@ meas2)[[1]]]