code.py

import numpy as np
import os
import subprocess
import ctypes
import tempfile
import math

def gc(L, n):
    """代码生成版本 - 极致优化"""
    
    class E: __slots__ = ()
    class C(E):
        __slots__ = ('v',)
        def __init__(s,v): s.v=float(v)
        def d(s,x): return C(0)
        def S(s): return s
        def c(s): return f"{int(s.v)}.0" if s.v==int(s.v) else str(s.v)
    class V(E):
        __slots__ = ('n',)
        def __init__(s,n): s.n=n
        def d(s,x): return C(1) if x==s.n else C(0)
        def S(s): return s
        def c(s): return s.n
    class A(E):
        __slots__ = ('a','b')
        def __init__(s,a,b): s.a,s.b=(a if isinstance(a,E)else C(a)),(b if isinstance(b,E)else C(b))
        def d(s,x): return A(s.a.d(x),s.b.d(x))
        def S(s):
            a,b=s.a.S(),s.b.S()
            if isinstance(a,C)and a.v==0:return b
            if isinstance(b,C)and b.v==0:return a
            if isinstance(a,C)and isinstance(b,C):return C(a.v+b.v)
            if isinstance(b,N):return U(a,b.a).S()  # a + (-b) => a - b
            return A(a,b)
        def c(s): return f"({s.a.c()}+{s.b.c()})"
    class U(E):
        __slots__ = ('a','b')
        def __init__(s,a,b): s.a,s.b=(a if isinstance(a,E)else C(a)),(b if isinstance(b,E)else C(b))
        def d(s,x): return U(s.a.d(x),s.b.d(x))
        def S(s):
            a,b=s.a.S(),s.b.S()
            if isinstance(b,C)and b.v==0:return a
            if isinstance(a,C)and a.v==0:return N(b).S()
            if isinstance(a,C)and isinstance(b,C):return C(a.v-b.v)
            if isinstance(b,N):return A(a,b.a).S()  # a - (-b) => a + b
            return U(a,b)
        def c(s): return f"({s.a.c()}-{s.b.c()})"
    class M(E):
        __slots__ = ('a','b')
        def __init__(s,a,b): s.a,s.b=(a if isinstance(a,E)else C(a)),(b if isinstance(b,E)else C(b))
        def d(s,x): return A(M(s.a,s.b.d(x)),M(s.a.d(x),s.b))
        def S(s):
            a,b=s.a.S(),s.b.S()
            if isinstance(a,C)and a.v==0:return C(0)
            if isinstance(b,C)and b.v==0:return C(0)
            if isinstance(a,C)and a.v==1:return b
            if isinstance(b,C)and b.v==1:return a
            if isinstance(a,C)and a.v==-1:return N(b).S()
            if isinstance(b,C)and b.v==-1:return N(a).S()
            if isinstance(a,C)and isinstance(b,C):return C(a.v*b.v)
            if isinstance(a,N)and isinstance(b,N):return M(a.a,b.a).S()  # (-a)*(-b) => a*b
            if isinstance(a,N):return N(M(a.a,b)).S()  # (-a)*b => -(a*b)
            if isinstance(b,N):return N(M(a,b.a)).S()  # a*(-b) => -(a*b)
            return M(a,b)
        def c(s): return f"({s.a.c()}*{s.b.c()})"
    class N(E):
        __slots__ = ('a',)
        def __init__(s,a): s.a=a if isinstance(a,E)else C(a)
        def d(s,x): return N(s.a.d(x))
        def S(s):
            a=s.a.S()
            if isinstance(a,C):return C(-a.v)
            if isinstance(a,N):return a.a
            return N(a)
        def c(s): return f"(-{s.a.c()})"
    class SN(E):
        __slots__ = ('a',)
        def __init__(s,a): s.a=a if isinstance(a,E)else C(a)
        def d(s,x): return M(CS(s.a),s.a.d(x))
        def S(s): return SN(s.a.S())
        def c(s): return f"sin({s.a.c()})"
    class CS(E):
        __slots__ = ('a',)
        def __init__(s,a): s.a=a if isinstance(a,E)else C(a)
        def d(s,x): return N(M(SN(s.a),s.a.d(x)))
        def S(s): return CS(s.a.S())
        def c(s): return f"cos({s.a.c()})"
    
    class SV:
        __slots__ = ('e',)
        def __init__(s,e): s.e=e if isinstance(e,E)else C(e)
        def __add__(s,o): return SV(A(s.e,o.e if isinstance(o,SV)else o))
        def __radd__(s,o): return s.__add__(o)
        def __sub__(s,o): return SV(U(s.e,o.e if isinstance(o,SV)else o))
        def __rsub__(s,o): return SV(U(C(o)if isinstance(o,(int,float))else o,s.e))
        def __mul__(s,o): return SV(M(s.e,o.e if isinstance(o,SV)else o))
        def __rmul__(s,o): return s.__mul__(o)
        def __neg__(s): return SV(N(s.e))
        def __truediv__(s,o):
            if isinstance(o,(int,float)):return SV(M(s.e,C(1.0/o)))
            raise NotImplementedError()
    
    class SM:
        sin=staticmethod(lambda x:SV(SN(x.e))if isinstance(x,SV)else math.sin(x))
        cos=staticmethod(lambda x:SV(CS(x.e))if isinstance(x,SV)else math.cos(x))
        sum=staticmethod(lambda x:sum(x)if hasattr(x,'__iter__')else x)
    
    class SA:
        def __init__(s,p,n): s.i=[SV(V(f"{p}[{i}]"))for i in range(n)]
        def __getitem__(s,i): return s.i[i]
        def __len__(s): return len(s.i)
    
    q_s,qd_s=SA("q",n),SA("qdot",n)
    L_r=L(q_s,qd_s,SM())
    L_e=L_r.e if isinstance(L_r,SV)else C(L_r)
    
    F_e=[L_e.d(f"q[{i}]").S()for i in range(n)]
    M_e,C_e=[],[]
    for i in range(n):
        dL=L_e.d(f"qdot[{i}]")
        for j in range(n):
            M_e.append(dL.d(f"qdot[{j}]").S())
            C_e.append(dL.d(f"q[{j}]").S())
    
    sc={}
    def col(e,d=0):
        if d>100:return
        if isinstance(e,(SN,CS)):
            k=e.a.c()
            if k not in sc:sc[k]=len(sc)
        if hasattr(e,'a'):col(e.a,d+1)
        if hasattr(e,'b'):col(e.b,d+1)
    for e in F_e+M_e+C_e:col(e)
    
    def rep(e):
        if isinstance(e,SN):return f"_s{sc[e.a.c()]}"
        if isinstance(e,CS):return f"_c{sc[e.a.c()]}"
        if isinstance(e,(C,V)):return e.c()
        if isinstance(e,A):return f"({rep(e.a)}+{rep(e.b)})"
        if isinstance(e,U):return f"({rep(e.a)}-{rep(e.b)})"
        if isinstance(e,M):return f"({rep(e.a)}*{rep(e.b)})"
        if isinstance(e,N):return f"(-{rep(e.a)})"
        return e.c()
    
    code=f"#define _GNU_SOURCE\n#include<math.h>\n#include<string.h>\n#define N {n}\n"
    
    # 生成计算核心函数
    code+="""
__attribute__((hot)) void cal_impl(double* __restrict__ q,double* __restrict__ qdot,double* __restrict__ qddot){
    double F[N],MM[N*N],y[N];
"""
    # 首先读取所有输入到局部变量
    for i in range(n):
        code+=f"double q{i}=q[{i}];\n"
    for i in range(n):
        code+=f"double qd{i}=qdot[{i}];\n"
    # sincos 调用
    for k,i in sorted(sc.items(),key=lambda x:x[1]):
        k_local = k
        for j in range(n):
            k_local = k_local.replace(f"q[{j}]", f"q{j}")
        code+=f"double _s{i},_c{i};sincos({k_local},&_s{i},&_c{i});\n"
    
    def rep_local(e):
        if isinstance(e,SN):return f"_s{sc[e.a.c()]}"
        if isinstance(e,CS):return f"_c{sc[e.a.c()]}"
        if isinstance(e,V):
            nm = e.n
            if nm.startswith("q["):
                idx = nm[2:-1]
                return f"q{idx}"
            elif nm.startswith("qdot["):
                idx = nm[5:-1]
                return f"qd{idx}"
            return nm
        if isinstance(e,C):return e.c()
        if isinstance(e,A):return f"({rep_local(e.a)}+{rep_local(e.b)})"
        if isinstance(e,U):return f"({rep_local(e.a)}-{rep_local(e.b)})"
        if isinstance(e,M):return f"({rep_local(e.a)}*{rep_local(e.b)})"
        if isinstance(e,N):return f"(-{rep_local(e.a)})"
        return e.c()
    
    for i in range(n):
        code+=f"F[{i}]={rep_local(F_e[i])};\n"
    
    # 只输出非零的 M 矩阵元素
    for i in range(n):
        for j in range(n):
            e = M_e[i*n+j]
            code+=f"MM[{i*n+j}]={rep_local(e)};\n"
    
    # 计算 b = F - C*qdot (直接计算，不存储 C)
    for i in range(n):
        terms = [f"F[{i}]"]
        for j in range(n):
            e = C_e[i*n+j]
            if not (isinstance(e, C) and e.v == 0):
                terms.append(f"({rep_local(e)})*qd{j}")
        code += f"y[{i}]={'-'.join(terms)};\n"
    
    # 内联求解器 - 优化版本
    code+="""
    // Gaussian elimination with partial pivoting - optimized
    for(int k=0;k<N;k++){
        int m=k;
        double mv=fabs(MM[k*N+k]);
        for(int i=k+1;i<N;i++){double v=fabs(MM[i*N+k]);if(v>mv){mv=v;m=i;}}
        if(m!=k){
            double*pk=MM+k*N,*pm=MM+m*N;
            for(int j=k;j<N;j++){double t=pk[j];pk[j]=pm[j];pm[j]=t;}
            double t=y[k];y[k]=y[m];y[m]=t;
        }
        double pivot=MM[k*N+k];
        for(int i=k+1;i<N;i++){
            double f=MM[i*N+k]/pivot;
            MM[i*N+k]=0;
            for(int j=k+1;j<N;j++)MM[i*N+j]-=f*MM[k*N+j];
            y[i]-=f*y[k];
        }
    }
    for(int i=N-1;i>=0;i--){
        double s=y[i];
        for(int j=i+1;j<N;j++)s-=MM[i*N+j]*qddot[j];
        qddot[i]=s/MM[i*N+i];
    }
}"""
    
    td=tempfile.mkdtemp()
    cf,sf=os.path.join(td,"c.c"),os.path.join(td,"c.so")
    with open(cf,"w")as f:f.write(code)
    r=subprocess.run(["gcc","-Ofast","-march=native","-funroll-loops","-ftree-vectorize","-fno-plt","-flto","-shared","-fPIC",cf,"-o",sf,"-lm"],capture_output=True)
    if r.returncode!=0:raise RuntimeError(r.stderr.decode())
    
    lib=ctypes.CDLL(sf)
    _f=lib.cal_impl
    _f.argtypes=[ctypes.c_void_p]*3
    _f.restype=None
    
    # 使用 from_buffer 获取指针，并完全缓存
    c_double_n = ctypes.c_double * n
    _addressof = ctypes.addressof
    _from_buffer = c_double_n.from_buffer
    
    # 缓存 from_buffer 结果和 addressof 结果
    _last_q = _last_qd = _last_qdd = None
    _c_q = _c_qd = _c_qdd = None
    _ptr_q = _ptr_qd = _ptr_qdd = 0
    
    def cal(q,qdot,qddot):
        nonlocal _last_q, _last_qd, _last_qdd, _c_q, _c_qd, _c_qdd, _ptr_q, _ptr_qd, _ptr_qdd
        if q is not _last_q:
            _last_q = q
            _c_q = _from_buffer(q)
            _ptr_q = _addressof(_c_q)
        if qdot is not _last_qd:
            _last_qd = qdot
            _c_qd = _from_buffer(qdot)
            _ptr_qd = _addressof(_c_qd)
        if qddot is not _last_qdd:
            _last_qdd = qddot
            _c_qdd = _from_buffer(qddot)
            _ptr_qdd = _addressof(_c_qdd)
        _f(_ptr_q, _ptr_qd, _ptr_qdd)
    return cal