The below code is one approach.
import numpy as np
Polynomial = np.polynomial.Polynomial
T = 2*60 + 15.2 # time to end of stage-one burn (s)
# acceleration, speed and distance polynomials
A = Polynomial((2.198, 2.842e-2, 1.061e-3)) # m.s-2
V = A.integ() # speed (m.s-1)
D = V.integ() # distance travelled (m)
print('Distance travelled after {:.1f} s = {:.1f} km'.format(T, D(T)/1.e3))
# average pitch angle of the rocket(from vertical)
theta = np.radians(12.)
# Ground temperature (K) and mean troposphere lapse rate (K.m-1)
# T(z) = T0 + Gamma.z, where z is the altitude in m
T0, Gamma = 302., 6. / 1000
# gamma = Cp/Cv and mean molar mass of atmosphere (kg.mol-1)
gamma, M = 1.4, 0.0288
# The gas constant (J.K-1.mol-1)
R = 8.314
# A lambda function to calculate the speed of sound, c(z) (m.s-1)
c = lambda z: np.sqrt(gamma * R * (T0 - Gamma*z)/M)
# a time grid across the stage-one burn
t = np.linspace(0, T, 1000)
# find the index of the point at which V > c
for i_mach1, t_mach1 in enumerate(t):
if V(t_mach1) > c(D(t_mach1)):
break
# altitude at Mach 1
z_mach1 = D(t_mach1) * np.cos(theta)
print('Mach 1 achieved at z = {:.1f} km, t = {:.1f} s, '
'where the local speed of sound is c = {:.1f} m.s-1'
.format(z_mach1 / 1000, t_mach1, c(z_mach1)))
The speed and distance travelled by the rocket are calculated by direct integration of the acceleration polynomial, using the default boundary conditions V(0)=0 and D(0)=0.
The speed of sound is calculated from the provided formula using the provided lapse rate, $\Gamma$:
$$
c = \sqrt{\frac{\gamma RT}{M}} = \sqrt{\frac{\gamma R(T_0 - \Gamma z)}{M}}
$$
c is defined as a lambda function using NumPy methods, so that it can return an array of values when passed an array of altitudes, z.
Next we set up a grid of 1000 time points and find the time and index of the first time the rocket's speed exceeds the speed of sound. There are several approaches to doing this, including numerically finding the roots of the equation V(t) - c(D(t)), but we choose to simply iterate over the time array and bail when this condition is met. The corresponding altitude can be calculated from the distance travelled by straightforward geometry using the average angle of the rocket from the vertical.