@ti.func
def sinh(x):
    return ti.exp(x) / 2 - ti.exp(-x) / 2

@ti.func
def project_strain(F_input, alpha_input):
    # Copy input variables to local variables
    F = F_input
    logJp = alpha_input

    # Perform Singular Value Decomposition (SVD) on F
    U, sig, V = ti.svd(F)
    sigma = ti.Vector([sig[0, 0], sig[1, 1]])  # Extract singular values into a vector

    # Compute determinant J
    J = sigma[0] * sigma[1]

    # Calculate p0 based on hardening
    p0 = kappa * (1e-5 + sinh(xi * ti.max(-logJp, 0.0)))
    p0 = ti.max(p0, 0.0)  # Ensure p0 is non-negative

    # Compute trial deviatoric stress
    B_hat_trial = sigma * sigma
    B_hat_trial_avg = B_hat_trial.sum() / dim
    s_hat_trial = mu * J**(-2.0 / dim) * (B_hat_trial - B_hat_trial_avg)

    # Compute trial pressure
    prime = kappa / 2.0 * (J - 1.0 / J)
    p_trial = -prime * J

    # Compute yield function value y
    y_s_half_coeff = (6.0 - dim) / 2.0 * (1.0 + 2.0 * beta)
    y_p_half = M * M * (p_trial + beta * p0) * (p_trial - p0)
    y = y_s_half_coeff * s_hat_trial.norm_sqr() + y_p_half

    # Project sigma if p_trial exceeds bounds
    p_min = beta * p0
    updated = False  # Flag to check if we need to update F and alpha

    tolerance = 1e-6  # Define a small tolerance for numerical errors

    if p_trial > p0 + tolerance:
        # Overflow in compression
        Je_new_argument = -2.0 * p0 / kappa + 1.0
        Je_new_argument = ti.max(Je_new_argument, 0.0)  # Ensure non-negative argument
        Je_new = ti.sqrt(Je_new_argument)
        sigma_new = ti.Vector([Je_new ** (1.0 / dim)] * dim)
        F = U @ ti.Matrix([[sigma_new[0], 0.], [0., sigma_new[1]]]) @ V.transpose()
        if hardening_on:
            logJp += ti.log(J / Je_new)
        updated = True

    elif p_trial < -p_min - tolerance:
        # Overflow in tension
        Je_new_argument = 2.0 * p_min / kappa + 1.0
        Je_new_argument = ti.max(Je_new_argument, 0.0)  # Ensure non-negative argument
        Je_new = ti.sqrt(Je_new_argument)
        sigma_new = ti.Vector([Je_new ** (1.0 / dim)] * dim)
        F = U @ ti.Matrix([[sigma_new[0], 0.], [0., sigma_new[1]]]) @ V.transpose()
        if hardening_on:
            logJp += ti.log(J / Je_new)
        updated = True

    elif y < 1e-4:
        # Inside the yield surface; no update needed
        pass
    else:
        # Project onto yield surface
        s_hat_trial_norm = s_hat_trial.norm()
        factor = ti.sqrt(-y_p_half / y_s_half_coeff) / s_hat_trial_norm
        B_hat_new = J**(2.0 / dim) / mu * factor * s_hat_trial + B_hat_trial_avg

        # Ensure B_hat_new is valid
        B_hat_new = ti.max(B_hat_new, 0.0)

        # Update sigma and F
        sigma_new = ti.sqrt(B_hat_new)
        F = U @ ti.Matrix([[sigma_new[0], 0.], [0., sigma_new[1]]]) @ V.transpose()

        # Hardening update if necessary
        if (p0 > 1e-4 and p_trial < p0 - 1e-4 and p_trial > 1e-4 - p_min):
            p_center = (p0 - p_min) * 0.5
            q_trial = ti.sqrt((6.0 - dim) / 2.0) * s_hat_trial_norm
            direction = ti.Vector([p_center - p_trial, -q_trial]).normalized()

            # Solve quadratic for lambda
            C = M * M * (p_center + beta * p0) * (p_center - p0)
            B_quad = M * M * direction[0] * (2.0 * p_center - p0 + beta * p0)
            A_quad = M * M * direction[0]**2 + (1.0 + 2.0 * beta) * direction[1]**2

            discriminant = B_quad * B_quad - 4.0 * A_quad * C
            discriminant = ti.max(discriminant, 0.0)  # Ensure non-negative
            if discriminant >= 0:
                sqrt_discriminant = ti.sqrt(discriminant)
                l1 = (-B_quad + sqrt_discriminant) / (2.0 * A_quad)
                l2 = (-B_quad - sqrt_discriminant) / (2.0 * A_quad)

                # Choose the correct lambda
                lambda_candidate = l1 if (p_trial - p_center) * (p_center + l1 * direction[0] - p_center) > 0 else l2
                p_fake = p_center + lambda_candidate * direction[0]
                Je_new_fake_argument = abs(-2.0 * p_fake / kappa + 1.0)
                Je_new_fake_argument = ti.max(Je_new_fake_argument, 0.0)
                Je_new_fake = ti.sqrt(Je_new_fake_argument)
                if Je_new_fake > 1e-4 and hardening_on:
                    logJp += ti.log(J / Je_new_fake)
                updated = True

    # Return updated F and alpha
    alpha_updated = logJp if updated and hardening_on else alpha_input
    return F, alpha_updated

Parts of material collapse, fix that