Commit 60e65c45 authored by Zhang Xin's avatar Zhang Xin
Browse files

add mpi; add pointing quaternion; modify trans method

parent 329533c1

survey_sim/config/infoOutput.py

+5 −5
Original line number Diff line number Diff line
@@ -2,7 +2,7 @@ 
Author: Zhang Xin zhangx@bao.ac.cn

Date: 2024-11-08 15:12:55

LastEditors: Zhang Xin zhangx@bao.ac.cn

LastEditTime: 2025-01-14 09:38:41

LastEditTime: 2025-05-13 02:09:07

FilePath: /CSST_Survey/survey_sim/config/infooutput.py

Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE

'''
@@ -26,8 +26,8 @@ class InfoOutput(object): 
        fh.setFormatter(formatter)

        self.logger.addHandler(fh)



        hdr1 = "# JDTime lonitude(ecliptic) latitude(ecliptic) RA Dec sun_x sun_y sun_z moon_x moon_y moon_z sat_x sat_y sat_z sat_vel_x sat_vel_y sat_vel_z isInDeep area_wide area_deep isInSunSide expTime transAngle"

        fmt1 = "%15.8f %15.8f %15.8f %15.8f %15.8f %15.8f %15.8f %15.8f %15.8f %15.8f %15.8f %15.8f %15.8f %15.8f %15.8f %15.8f %15.8f %4d  %15.8f  %15.8f  %15.8f  %15.8f  %15.8f"

        hdr1 = "# JDTime lonitude(ecliptic) latitude(ecliptic) RA Dec sun_x sun_y sun_z moon_x moon_y moon_z sat_x sat_y sat_z sat_vel_x sat_vel_y sat_vel_z isInDeep area_wide area_deep isInSunSide expTime transAngle cmg energy"

        fmt1 = "%15.8f %15.8f %15.8f %15.8f %15.8f %15.8f %15.8f %15.8f %15.8f %15.8f %15.8f %15.8f %15.8f %15.8f %15.8f %15.8f %15.8f %4d  %15.8f  %15.8f  %15.8f  %15.8f  %15.8f  %15.8f  %15.8f"

        self.hdr = hdr1

        self.fmt = fmt1
@@ -57,10 +57,10 @@ class InfoOutput(object): 
        self.outInfo.close()

        # self.logger.close()



    def outInfo_add_obj(self, jdTime=2459766., p_lon_ecl=0., p_lat_ecl=0., p_ra=0, p_dec=0., sun=[0, 0, 0], moon=[0, 0, 0], sat=[0, 0, 0], sat_vel=[0, 0, 0], isInDeep=0, areaW=0.0, areaD=0.0, isInSunSide=1, exp_time=150., trans_Angle=1.):

    def outInfo_add_obj(self, jdTime=2459766., p_lon_ecl=0., p_lat_ecl=0., p_ra=0, p_dec=0., sun=[0, 0, 0], moon=[0, 0, 0], sat=[0, 0, 0], sat_vel=[0, 0, 0], isInDeep=0, areaW=0.0, areaD=0.0, isInSunSide=1, exp_time=150., trans_Angle=1., cmg=0., energy=0.):



        line = self.fmt % (

            jdTime, p_lon_ecl, p_lat_ecl, p_ra, p_dec, sun[0], sun[1], sun[2], moon[0], moon[1], moon[2], sat[0], sat[1], sat[2], sat_vel[0], sat_vel[1], sat_vel[2], isInDeep, areaW, areaD, isInSunSide, exp_time, trans_Angle)

            jdTime, p_lon_ecl, p_lat_ecl, p_ra, p_dec, sun[0], sun[1], sun[2], moon[0], moon[1], moon[2], sat[0], sat[1], sat[2], sat_vel[0], sat_vel[1], sat_vel[2], isInDeep, areaW, areaD, isInSunSide, exp_time, trans_Angle, cmg, energy)

        # if not line.endswith("\n"):

        line += "\n"

        self.outInfo.write(line)

survey_sim/constraints/_utils.py

+162 −1
Original line number Diff line number Diff line
@@ -298,6 +298,164 @@ def calculate_trace(matrix): 
    return trace_sum





def quaternion_from_axis_angle(axis, angle_rad):

    """

    给定旋转轴和角度，生成对应的单位四元数

    """

    axis = np.array(axis, dtype=float)

    axis = axis / np.linalg.norm(axis)

    half_angle = angle_rad / 2.0

    q0 = np.cos(half_angle)

    q_xyz = axis * np.sin(half_angle)

    return np.concatenate(([q0], q_xyz))





def quaternion_multiply(q1, q2):

    """

    计算两个四元数的乘积 q = q1 * q2

    四元数格式为 [q0, q1, q2, q3]

    """

    w1, x1, y1, z1 = q1

    w2, x2, y2, z2 = q2

    return np.array([

        w1*w2 - x1*x2 - y1*y2 - z1*z2,

        w1*x2 + x1*w2 + y1*z2 - z1*y2,

        w1*y2 - x1*z2 + y1*w2 + z1*x2,

        w1*z2 + x1*y2 - y1*x2 + z1*w2

    ])





def quaternion_conjugate(q):

    """

    计算两个四元数的共轭

    """

    q0, q1, q2, q3 = q

    return np.array([q0, -q1, -q2, -q3])





def rotate_vector_by_quaternion(v, q):

    """

    使用单位四元数 q 旋转向量 v



    参数:

        v: 3维向量 [vx, vy, vz]

        q: 单位四元数 [q0, q1, q2, q3]



    返回:

        v_rotated: 旋转后的3维向量

    """

    v_q = np.concatenate(([0.0], v))

    q_conj = quaternion_conjugate(q)

    qv = quaternion_multiply(q, v_q)

    qvq = quaternion_multiply(qv, q_conj)

    return qvq[1:]  # 只取向量部分





def axis_angle_from_quaternion(q):

    """

    将单位四元数转换为旋转轴和旋转角（弧度）



    参数:

        q: 四元数 [q0, q1, q2, q3]



    返回:

        axis: 旋转轴（单位向量）

        angle: 旋转角（单位：弧度）

    """

    q0, q1, q2, q3 = q

    q0 = np.clip(q0, -1.0, 1.0)

    angle = 2 * np.arccos(q0)



    sin_half_angle = np.sqrt(1 - q0*q0)



    if sin_half_angle < 1e-8:

        # 接近0度旋转，轴任意，这里返回Z轴

        axis = np.array([0.0, 0.0, 1.0])

    else:

        axis = np.array([q1, q2, q3]) / sin_half_angle



    return axis, angle





def transferlonlat2cardisian(lon_rad=0., lat_rad=0.):

    return [math.cos(lat_rad) * math.cos(

        lon_rad), math.cos(lat_rad) * math.sin(lon_rad), math.sin(lat_rad)]





def transfer2lonlat_quaternion(lon=90., lat=90., pa=0.):

    lon_rad = math.radians(lon)

    lat_rad = math.radians(lat)

    pa_rad = math.radians(pa)



    # q_x = quaternion_from_axis_angle([1, 0, 0], pa_rad)



    q_z = quaternion_from_axis_angle([0, 0, 1], lon_rad)

    n_y_vect = rotate_vector_by_quaternion([0, 1, 0], q_z)



    q_n = quaternion_from_axis_angle(n_y_vect, -1*lat_rad)



    q_t = quaternion_from_axis_angle([math.cos(lat_rad) * math.cos(

        lon_rad), math.cos(lat_rad) * math.sin(lon_rad), math.sin(lat_rad)], pa_rad)



    q_total = quaternion_multiply(q_t, q_n)

    q_total = quaternion_multiply(q_total, q_z)



    # q_total = quaternion_multiply(q_x, q_z)

    # q_total = quaternion_multiply(q_total, q_n)

    return q_total





def transfer2lonlat_quaternion_noPa(lon=90., lat=90.):

    lon_rad = math.radians(lon)

    lat_rad = math.radians(lat)

    q_z = quaternion_from_axis_angle([0, 0, 1], lon_rad)

    n_y_vect = rotate_vector_by_quaternion([0, 1, 0], q_z)



    q_n = quaternion_from_axis_angle(n_y_vect, -1*lat_rad)



    q_total = quaternion_multiply(q_n, q_z)



    return q_total





def rotate_quaternion_byself(lon=90., lat=90., pa=0., quaternion=[1, 0, 0, 0]):

    lon_rad = math.radians(lon)

    lat_rad = math.radians(lat)

    pa_rad = math.radians(pa)

    q_t = quaternion_from_axis_angle([math.cos(lat_rad) * math.cos(

        lon_rad), math.cos(lat_rad) * math.sin(lon_rad), math.sin(lat_rad)], pa_rad)

    q_total = quaternion_multiply(q_t, quaternion)

    return q_total





def get_RotationAngle_twoQuaternion(q_old=[1., 0., 0., 0.], q_new=[1., 0., 0., 0.]):

    q_old_conj = quaternion_conjugate(q_old)

    q_final = quaternion_multiply(q_new, q_old_conj)



    axis_final, angle_final = axis_angle_from_quaternion(q_final)

    if angle_final > math.pi:

        angle_final = 2*math.pi-angle_final

        axis_final = -1*axis_final

        q_final[1:] = -1*q_final[1:]

    return math.degrees(angle_final), axis_final, q_final





def get_RotationAngle_Quaternion(old_lon_lat, new_lon_lat, old_pa=0.0, new_pa=0.0):

    q_old = transfer2lonlat_quaternion(

        lon=old_lon_lat[0], lat=old_lon_lat[1], pa=old_pa)

    q_new = transfer2lonlat_quaternion(

        lon=new_lon_lat[0], lat=new_lon_lat[1], pa=new_pa)

    q_old_conj = quaternion_conjugate(q_old)

    q_final = quaternion_multiply(q_new, q_old_conj)



    axis_final, angle_final = axis_angle_from_quaternion(q_final)

    if angle_final > math.pi:

        angle_final = 2*math.pi-angle_final

        axis_final = -1*axis_final

        q_final[1:] = -1*q_final[1:]

    return math.degrees(angle_final), axis_final, q_final





@jit

def get_RotationAngleFromMatrix(mat):

    # assert np.allclose(np.dot(mat.T, mat), np.eye(3)), "R must be orthogonal"
@@ -559,6 +717,7 @@ def calculateTransTime(transAngle=1.0, surveyCons=None): 
    #     )

    #     tTime = angleVStime_i(transAngle)

    # print(tTime)

    tTime = 0

    if transAngle < angleVStime[0, 0]:

        tTime = 70

    elif transAngle == angleVStime[0, 0]:
@@ -570,7 +729,9 @@ def calculateTransTime(transAngle=1.0, surveyCons=None): 
                    angleVStime[1, i] * ((transAngle - angleVStime[0, i-1])) / \

                    (((angleVStime[0, i] - angleVStime[0, i-1])))

                break



    if transAngle > angleVStime[0, -1]:

        print("error trans angle:", transAngle)

        tTime = angleVStime[1, -1]

    return tTime + surveyCons.SHUTTER_TIME * 2.0



    # int i = 0;

survey_sim/constraints/beta_time_constraint.py

+2 −0
Original line number Diff line number Diff line
@@ -4,6 +4,7 @@ import numpy as np 
import math

from survey_sim.ephemeris import locate_sun

import matplotlib.pyplot as plt

import gc





# def get_betaAngle(time = 2459799, orbitData = None):
@@ -171,6 +172,7 @@ class beta_time_constraint(object): 
        self.beta_time_seg = beta_time_seg_d_flat.reshape(

            beta_time_seg.shape[0] - d_ids[0].shape[0], 2

        )

        gc.collect()



        # return beta_time_seg_d

survey_sim/constraints/obscure_constraint.py

+53 −0
Original line number Diff line number Diff line
@@ -39,6 +39,24 @@ def isObscureBySun(sun=None, p=None, constr=None): 
    return sunObscurFlag





def isObscureBySun_array(sun=None, p=None, constr=None):



    modSun = np.linalg.norm(sun)



    modP = np.linalg.norm(p, axis=1)

    # modP = p[0] * p[0] + p[1] * p[1] + p[2] * p[2]

    # modSun = sun[0] * sun[0] + sun[1] * sun[1] + sun[2] * sun[2]

    cosA = (sun[0] * p[:, 0] + sun[1] * p[:, 1] + sun[2]

            * p[:, 2]) / modSun * modP



    return cosA <= constr.sun_los_angle_cos

    # sunObscurFlag = 2

    # if cosA <= constr.sun_los_angle_cos:

    #     sunObscurFlag = cosA



    # return sunObscurFlag





"""

description: 是否被月球遮挡

param {*} moon: 月球位置，笛卡尔坐标
@@ -62,6 +80,23 @@ def IsObscureByMoon(moon=None, p=None, constr=None): 
    return moonObscureFlag





def IsObscureByMoon_array(moon=None, p=None, constr=None):



    modMoon = np.linalg.norm(moon)



    modP = np.linalg.norm(p, axis=1)



    cosA = (moon[0] * p[:, 0] + moon[1] * p[:, 1] +

            moon[2] * p[:, 2]) / modMoon * modP



    return cosA <= constr.moon_los_angle_cos

    # moonObscureFlag = 2

    # if cosA <= constr.moon_los_angle_cos:

    #     moonObscureFlag = cosA



    # return moonObscureFlag





"""

description: 

return {*} isObscure 1:被遮挡 0：未被遮挡
@@ -238,6 +273,8 @@ def aquireShadowTime(curTime=2459766.0, orbitDat=None, ephlib=None): 


    return [shadow_start, shadow_end]



# 被遮挡设为True





def IsObscureByEarth_firstCut(sat=None, p=None):

    modSat = np.linalg.norm(sat)
@@ -252,3 +289,19 @@ def IsObscureByEarth_firstCut(sat=None, p=None): 
        return -1

    else:

        return 1



# 被遮挡设为False





def IsObscureByEarth_firstCut_array(sat=None, p=None):



    num = len(p)

    modSat = np.linalg.norm(sat)



    modPoint = np.linalg.norm(p, axis=1)



    withLocalZenithAngle = (p[:, 0] * sat[0] + p[:, 1] * sat[1] + p[:, 2] * sat[2]) / (

        modPoint * modSat

    )



    return withLocalZenithAngle >= 0

survey_sim/constraints/surveyConstraint.py

+5 −1
Original line number Diff line number Diff line
@@ -2,7 +2,7 @@ 
Author: Zhang Xin zhangx@bao.ac.cn

Date: 2020-06-17 17:03:15

LastEditors: Zhang Xin zhangx@bao.ac.cn

LastEditTime: 2024-11-14 08:29:09

LastEditTime: 2025-04-08 14:20:26

FilePath: /CSST_Survey/survey_sim/constraints/surveyConstraint.py

Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE

'''
@@ -76,3 +76,7 @@ class surveyConstraint(object): 


        self.DEC60_PRIOR_TIME = 11

        self.HIGH_LATITUDE_PRIOR_TIME = 10.5



    def _get_arrtr_(self, arrt="sun_los_angle_cos"):

        if arrt == "sun_los_angle_cos":

            return self.sun_los_angle_cos