def _decoherence_matrix(
    cool_rate: float, dephase_rate: float, heat_rate: float = 0.0, dim: int = 2
) -> np.ndarray:
    """Construct a rate matrix associated with decay and dephasing.

    The units of the matrix match the units of the rates specified.
    This matrix can be used to construct a noise channel after rescaling
    by an idling time (to make it dimensionless).

    Args:
        cool_rate: Decay rate of the system, usually 1 / T_1
        dephase_rate: Static dephasing rate of the system, usually 1 / T_phi
        heat_rate: Heating rate of the system (default 0).
        dim: Number of energy levels to include (default 2).

    Returns:
        np.ndarray rate matrix for decay and dephasing.
    """
    # Heating (related to a^dag)
    rate_matrix = np.diag(np.arange(1, dim) * heat_rate, 1).T.astype(float)
    # Cooling (related to a)
    rate_matrix += np.diag(np.arange(1, dim) * cool_rate, 1)
    # Dephasing (related to n=a^dag * a)
    # We specify i^2 since we take square root to get the Lindblad op later.
    rate_matrix += np.diag(dephase_rate * np.arange(dim) ** 2)
    return rate_matrix
Construct a rate matrix associated with decay and dephasing.

    The units of the matrix match the units of the rates specified.
    This matrix can be used to construct a noise channel after rescaling
    by an idling time (to make it dimensionless).

    Args:
        cool_rate: Decay rate of the system, usually 1 / T_1
        dephase_rate: Static dephasing rate of the system, usually 1 / T_phi
        heat_rate: Heating rate of the system (default 0).
        dim: Number of energy levels to include (default 2).

    Returns:
        np.ndarray rate matrix for decay and dephasing.
    
cirq-core/cirq/devices/thermal_noise_model.py