Advent Day 24 Part 2

import re
import sys
import numpy as np

def advent_day24_part2(filestr):
    field = [x for x in filestr.split('\n') if x]
    nrows, ncols = len(field), len(field[0])

    stones = []
    for line in field:
        m = re.split(r'[, @]+', line)
        if m:
            vals = [int(x) for x in m]
            p = tuple(vals[:3])
            v = tuple(vals[3:])
            stones.append((p, v))

    # x = [sx sy sz wx wy wz t1 t2 ... tn]
    # x_n+1 = x_n - J(x_n)^-1 f(x_n)
    # f(x) = sx + wx * t1 - (px1 + vx1 * t1)
    #      = (sx - px1) + (wx - vx1) * t1
    # ... and again for y and z
    # ... and again for all n
    # J_x = [ 1  0  0 t1  0  0  wx-vx1  0  0 ...]
    # J_y = [ 0  1  0  0 t1  0  wy-vy1  0  0 ...]
    # J_z = [ 0  0  1  0  0 t1  wz-vz1  0  0 ...]

    ps = np.array(tuple(row[0] for row in stones), dtype=np.float64)
    vs = np.array(tuple(row[1] for row in stones), dtype=np.float64)
    x = np.array([0, 0, 0, 1, 1, 1] + [0] * nrows, dtype=np.float64)

    def jacobian(x, vs):
        data = np.zeros((nrows * 3, 6 + nrows))
        for i in range(nrows):
            for j in range(3):
                r = 3 * i + j
                data[r, j] = 1
                data[r, 3 + j] = x[6 + i]
                data[r, 6 + i] = x[3 + j] - vs[i][j]
        return data

    def func(x, ps, vs):
        data = np.zeros(nrows * 3)
        for i in range(nrows):
            for j in range(3):
                r = 3 * i + j
                t = x[6 + i]
                sp = x[j] + x[3 + j] * t
                hp = ps[i][j] + vs[i][j] * t
                data[r] = sp - hp
        return data

    sum0 = 0
    for i in range(100):
        print(i, np.sum(x[:3]), np.int64(x[0:3]), x[3:6])
        j = jacobian(x, vs)
        f = func(x, ps, vs)
        dx, resid, rank, sing = np.linalg.lstsq(j, f, rcond=None)
        x -= dx
        sum1 = np.sum(x[:3])
        if abs(sum1 - sum0) < 0.1:
            break
        sum0 = sum1
    print(np.int64(x[:3]), np.int64(x[3:6]))
    print(int(sum0))

def readfile(file_path):
    with open(file_path, 'r') as file:
        return file.read()

if __name__ == '__main__':
    if len(sys.argv) >= 2:
        filestr = readfile(sys.argv[1])
    else:
        filestr = ''
    advent_day24_part2(filestr)