Suppose the quantum system has state $$a|0_A\rangle|0_B\rangle + b|1_A\rangle|1_B\rangle.$$ I want to prove that system $A$ has reduced density matrix of $a^2|0_A\rangle\langle 0_A| + b^2 |1_A\rangle\langle 1_A|$ when $0_A\rangle$ and $|1_A\rangle$ are orthogonal (and assume that $0_B\rangle$ and $1_B\rangle$ are orthogonal), but I am having a hard time mathematizing calculations.
So how is the trace over $B$ mathematically done?
The partial trace $\mathrm{Tr}_B$ is defined as:
Note that the second half is not enough to specify what happens to arbitrary operators, but since you can decompose any arbitrary (nice-enough) operator on $\mathcal H_A\otimes \mathcal H_B$ as a sum of tensor-product terms, then the requirement that the partial trace be linear completely fixes its behaviour.
To calculate its action on an operator like \begin{align} \rho & = |\psi\rangle\langle\psi| \\ & = |a|^2 |0_A 0_B\rangle \langle 0_A 0_B| + a b^* |0_A 0_B\rangle \langle 1_A 1_B| \\ & \qquad + a^* b |1_A 1_B\rangle \langle 0_A 0_B| + |b|^2 |1_A 1_B\rangle \langle 1_A 1_B|, \\ & = |a|^2 |0_A \rangle \langle 0_A|\otimes |0_B\rangle \langle 0_B| + a b^* |0_A \rangle \langle 1_A|\otimes |0_B\rangle \langle 1_B| \\ & \qquad + a^* b |1_A \rangle \langle 0_A|\otimes |1_B\rangle \langle 0_B| + |b|^2 |1_A \rangle \langle 1_A|\otimes |1_B\rangle \langle 1_B|, \end{align} you have to apply those two properties sequentially: first you break up the sum, \begin{align} \mathrm{Tr}_B(\rho) & = |a|^2 \,\mathrm{Tr}_B(|0_A \rangle \langle 0_A|\otimes |0_B\rangle \langle 0_B|) + a b^* \, \mathrm{Tr}_B(|0_A \rangle \langle 1_A|\otimes |0_B\rangle \langle 1_B| ) \\ & \qquad + a^* b \,\mathrm{Tr}_B(|1_A \rangle \langle 0_A|\otimes |1_B\rangle \langle 0_B| ) + |b|^2 \, \mathrm{Tr}_B(|1_A \rangle \langle 1_A|\otimes |1_B\rangle \langle 1_B|), \end{align} then you pass the traces to the right-hand side, \begin{align} \mathrm{Tr}_B(\rho) & = |a|^2 |0_A \rangle \langle 0_A|\,\mathrm{Tr}_B( |0_B\rangle \langle 0_B|) + a b^* |0_A \rangle \langle 1_A|\, \mathrm{Tr}_B(|0_B\rangle \langle 1_B| ) \\ & \qquad + a^* b |1_A \rangle \langle 0_A|\,\mathrm{Tr}_B( |1_B\rangle \langle 0_B| ) + |b|^2 |1_A \rangle \langle 1_A| \,\mathrm{Tr}_B( |1_B\rangle \langle 1_B|), \end{align} and then you calculate the traces: \begin{align} \mathrm{Tr}_B(\rho) & = |a|^2 |0_A \rangle \langle 0_A| \times 1 + a b^* |0_A \rangle \langle 1_A| \times 0 \\ & \qquad + a^* b |1_A \rangle \langle 0_A| \times 0 + |b|^2 |1_A \rangle \langle 1_A| \times 1 \\ & = |a|^2 |0_A \rangle \langle 0_A| + |b|^2 |1_A \rangle \langle 1_A|. \end{align} Easy!
As some additional practice, try calculating the partial trace of $\rho=|\psi⟩⟨\psi|$ as you go from an entangled state to a separable one, by taking, say, $$ \psi = a|0_A\rangle|0_B\rangle + b|1_A\rangle|1_B\rangle +\sin(\theta)\sqrt{ab}(|0_A\rangle|1_B\rangle + b|1_A\rangle|0_B\rangle), $$ which simplifies to your state at $\theta=0$ and to the product state $$(\sqrt{a}|0_A⟩+\sqrt{b}|1_A⟩) \otimes (\sqrt{a}|0_B⟩+\sqrt{b}|1_B⟩)$$ at $\theta=\pi/2$: you should see that as the parameter $\theta$ increases, the relevance of the off-diagonal elements of $\mathrm{Tr}_B(\rho)$ increases as well.
Let's assume that the states $|0\rangle,|1\rangle$ are normalized. Let's also assume $|a|^2+|b|^2 = 1$ so that $|\psi\rangle = a |0_A 0_B\rangle+ b |1_A 1_B\rangle$ is normalized (this is all done so that $\rho$ is properly normalized).
We have then $\rho = |\psi\rangle\langle\psi|= |a|^2 |0_A 0_B\rangle \langle 0_A 0_B| + a b^* |0_A 0_B\rangle \langle 1_A 1_B| + a^* b |1_A 1_B\rangle \langle 0_A 0_B| + |b|^2 |1_A 1_B\rangle \langle 1_A 1_B|$.
Taking the trace over $B$ we get $\rho_A = \langle 0_B| \rho|0_B\rangle + \langle 1_B| \rho|1_B\rangle = |a|^2 |0_A\rangle \langle 0_A| + |b|^2 |1_A\rangle \langle 1_A|$.
