How to derive reduced density matrix of $a|0_A\rangle|0_B\rangle + b|1_A\rangle|1_B\rangle$ 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?
 A: The partial trace $\mathrm{Tr}_B$ is defined as:


*

*a linear map $\mathrm{Tr}_B:\mathcal{L}(\mathcal H_A\otimes \mathcal H_B) \to \mathcal{L}(\mathcal H_A)$ that goes from the space $\mathcal{L}(\mathcal H_A\otimes \mathcal H_B)$ of linear operators on $\mathcal H_A\otimes \mathcal H_B$ to the space of linear operators on $\mathcal H_A$, such that

*when it is given a tensor-product operator $O_A\otimes O_B$, it returns the left half and takes the trace of the right half, i.e. $\mathrm{Tr}_B(O_A\otimes O_B) = O_A\,\mathrm{Tr}(O_B).$


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.
A: 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|$.
