Square root and logarithm for positive branch cut

# %% Square root and logarithm for positive branch cut

def logth(z, theta):
    argument = np.angle(z)
    modulus = np.absolute(z)
    argument = np.mod(argument + theta, 2 * np.pi) - theta
    return np.log(modulus) + 1j * argument

log1 = lambda z: logth(z, 2*np.pi)

def sqrtth(z, theta):
    argument = np.angle(z)
    modulus = np.abs(z)
    argument = np.where(argument == 0, 0, np.mod(argument + theta, 2 * np.pi) - theta)
    return np.sqrt(modulus) * np.exp(1j * argument / 2)

sqrt1 = lambda z: sqrtth(z, 2*np.pi)