読者です 読者をやめる 読者になる 読者になる

はやぶさ2のpythonプロットをリファクタリング

python matplot matplotlib Hayabusa2 Ryugu プログラミング

前回の記事で書いたプログラムを先日使う機会があったが、ものすごい使いづらかった。。。
ということで、空き時間でちまちまとリファクタリングした。

前回の記事のコードがなぜ使いづらいのかは、この部分から分かるはず…

class HAYA2:
    def __init__(self):
        self.lst = []
        h2txt = open("./haya2_orbit_jaxa.txt","r")
        h2list = [x.strip() for x in h2txt.readlines()][51:]
        for line in h2list:
            _d = {}
            # _d['date'] = line[0:10].replace("/", "") # target date
            _date = datetime.datetime.strptime(line[0:19], "%Y/%m/%d.%H:%M:%S")
            _d['date'] = datetime.date(_date.year, _date.month, _date.day) # date type
            _d['lp']   = int(line[22:26])          # L+ [days]
            _d['x']  = float(line[29:38])        # X pos. [au]
            _d['y']  = float(line[39:48])        # Y pos. [au]
            _d['z']  = float(line[49:58])        # Z pos. [au]
            _d['ex'] = float(line[59:68])
            _d['ey'] = float(line[69:78])
            _d['ez'] = float(line[79:88])
            _d['rs'] = float(line[122:129])      # distance of Sun-Haya2 [10**4 km]
            _d['re'] = float(line[133:140])      # distance of Earth-Haya2 [10**4 km]
            _d['ra'] = float(line[144:151])      # distance of 1999JU3-Haya2 [10**4 km]
            _d['vs']  = float(line[154:159])     # velocity of Haya2 on Sun [km/sec]
            _d['ve']  = float(line[162:167])     # velocity of Haya2 on Earth [km/sec]
            _d['alpha'] = float(line[169:176])   # ra [deg]
            _d['delta'] = float(line[179:185])   # dec [deg]
            _d['Dflt']  = float(line[186:194])   # distance of fling [10**4 km]

            _d['px'], _d['py'], _d['pz'] = convertCood(_d['x'], _d['y'], _d['z'])

            self.lst.append(_d)

def plotHaya2(day):
    haya2 = HAYA2()
    for h in haya2.lst[0:365]:
        plt.plot(h['px'], h['py'], "k.", ms=1)
    for h in haya2.lst[0:365:30]:
        plt.plot(h['px'], h['py'], "c.")
        plt.text(h['px'], h['py'], "${date:%m/%d}^{{\mathrm{{'}}{date:%y}}}$".format(**h),
                 ha='center', va='bottom', fontsize=6, color='blue')

せっかく、コンストラクタでデータ読み込みをしたのに、そのインスタンスを外からブン回す。
あまりにもひどいので、できるだけクラスメソッドとして整理して、以下のようなコードに

class JAXA(Planet):
    def __init__(self, name, data_file):
        self.name = name

        if self.name == "Hayabusa2":
            self.data = []
            h2list = [x.strip() for x in open(data_file,'r',encoding='utf-8').readlines()][51:]
            for line in h2list:
                _d = {}
                # _d['date'] = line[0:10].replace("/", "") # target date
                _date = datetime.datetime.strptime(line[0:19], "%Y/%m/%d.%H:%M:%S")
                _d['date'] = datetime.date(_date.year, _date.month, _date.day) # date type
                _d['lp']   = int(line[22:26])          # L+ [days]
                _d['x']  = float(line[29:38])        # X pos. [au]
                _d['y']  = float(line[39:48])        # Y pos. [au]
                _d['z']  = float(line[49:58])        # Z pos. [au]
                _d['ex'] = float(line[59:68])
                _d['ey'] = float(line[69:78])
                _d['ez'] = float(line[79:88])
                _d['rs'] = float(line[122:129])      # distance of Sun-Haya2 [10**4 km]
                _d['re'] = float(line[133:140])      # distance of Earth-Haya2 [10**4 km]
                _d['ra'] = float(line[144:151])      # distance of 1999JU3-Haya2 [10**4 km]
                _d['vs']  = float(line[154:159])     # velocity of Haya2 on Sun [km/sec]
                _d['ve']  = float(line[162:167])     # velocity of Haya2 on Earth [km/sec]
                _d['alpha'] = float(line[169:176])   # ra [deg]
                _d['delta'] = float(line[179:185])   # dec [deg]
                _d['Dflt']  = float(line[186:194])   # distance of fling [10**4 km]

                # _d['px'], _d['py'], _d['pz'] = convertCood(_d['x'], _d['y'], _d['z'])

                self.data.append(_d)

        if self.name == "Ryugu":
            self.data = []
            h2list = [x.strip() for x in open(data_file,'r',encoding='utf-8').readlines()][51:]
            for line in h2list:
                _d = {}
                # _d['date'] = line[0:10].replace("/", "") # target date
                _date = datetime.datetime.strptime(line[0:19], "%Y/%m/%d.%H:%M:%S")
                _d['date'] = datetime.date(_date.year, _date.month, _date.day) # date type
                _d['lp']   = int(line[22:26])          # L+ [days]
                _d['x']  = float(line[89:98])        # X pos. [au]
                _d['y']  = float(line[99:108])        # Y pos. [au]
                _d['z']  = float(line[109:118])        # Z pos. [au]
                _d['ra'] = float(line[144:151])      # distance of 1999JU3-Haya2 [10**4 km]

                # _d['px'], _d['py'], _d['pz'] = convertCood(_d['x'], _d['y'], _d['z'])

                self.data.append(_d)


    def drawOrbitJAXA(self, ax, params, begin_lp = None, end_lp = None):
        '''
        plot line of target planet (for planet orbit)
        '''
        begin_lp = 0 if begin_lp==None else begin_lp
        end_lp = 1279 if end_lp==None else end_lp

        xl, yl = [], []
        for d in self.data[begin_lp:end_lp]:
            px, py, pz = self.convertCood(d['x'], d['y'], d['z'], params)
            xl.append(px)
            yl.append(py)

        ax.plot(xl, yl, '--', lw=0.5,
                color = {'Hayabusa2' : '#00bfff',
                         'Ryugu': 'k'
                     }[self.name])

    def plotPointJAXA(self, ax, params, begin_lp, end_lp = None, days_interval = None):
        '''
        plot point of target planet on target date
        '''
        end_lp = begin_lp+1 if end_lp==None else end_lp
        days_interval = 1 if days_interval==None else days_interval

        for d in self.data[begin_lp:end_lp:days_interval]:
            px, py, pz = self.convertCood(d['x'], d['y'], d['z'], params)
            ax.plot(px, py,
                    {'Hayabusa2' : 'v',
                     'Ryugu': 'p'
                    }[self.name],
                    ms=12,
                    color = {'Hayabusa2' : '#00bfff',
                             'Ryugu': 'k'
                            }[self.name])

    def textDateJAXA(self, ax, params, begin_lp, end_lp = None, days_interval = None):
        '''
        caption text of planets
        '''
        end_lp = begin_lp+1 if end_lp==None else end_lp
        days_interval = 1 if days_interval==None else days_interval

        for d in self.data[begin_lp:end_lp:days_interval]:
            px, py, pz = self.convertCood(d['x'], d['y'], d['z'], params)
            ax.text(px, py,
                    "${0:%m/%d}^{{\mathrm{{'}}{0:%y}}}$".format(d['date']),
                    fontsize=6, ha='left', va='top')

    def textAngleEVEJAXA(self, ax, params, begin_lp, end_lp = None, days_interval = None):
        '''
        for exhibition function
        angle of Vernal Equinox day's Earth positon <-> Sun position <-> Planet position
        '''
        end_lp = begin_lp+1 if end_lp==None else end_lp
        days_interval = 1 if days_interval==None else days_interval

        for d in self.data[begin_lp:end_lp:days_interval]:
            px, py, pz = self.convertCood(d['x'], d['y'], d['z'], params)
            a = np.array([px, py])
            b = np.array(params['EVE'])
            ang = self.angle2vector(a, b) # ang[radian]
            ang_deg = math.degrees(ang)
            print("{:s} angle(x-y): {:.3f}".format(self.name, ang_deg))
            ax.text(px, py, "{:.1f}".format(ang_deg), fontsize=6, ha='left', va='bottom')


def main():
    params = {'inner': None,
              'theta': 6,
              'phi': -8,
              'mag': 0.245/1.,
              'EVE': [0, 0],
              'EVEday': datetime.date(2014, 3, 21)}

    target_date = datetime.date(2015, 12, 3) # plot position on target_date

    ####### INNER ########
    params['inner'] = True

    fig = plt.figure(figsize=(5,5))
    ax  = fig.add_subplot(111)
    ax.set_xlabel("$x[\mathrm{au}]$")
    ax.set_ylabel("$y[\mathrm{au}]$")
    ax.axis('equal')
    if params['inner']:
        ax.axis([-2,2,-2,2])
        ax.set_title('inner')
    else:
        ax.axis([-5,5,-5,5])
        ax.set_title('outer')

    ax.plot(0, 0, "ro") # SUN

    Earth = Planet("Earth")
    params['EVE'] = Earth.calcEVE(ax, params)
    Earth.plotPointOnEVE(ax, params) # plot Equinox of Earth J2000

    Mercury = Planet("Mercury")
    Venus = Planet("Venus")
    Mars = Planet("Mars")
    Jupiter = Planet("Jupiter")
    Saturn = Planet("Saturn")
    Uranus = Planet("Uranus")
    Neptune = Planet("Neptune")
    Pluto = Planet("Pluto")

    Mercury.drawOrbit(ax, params)
    Venus.drawOrbit(ax, params)
    Earth.drawOrbit(ax, params)
    Mars.drawOrbit(ax, params)
    # Jupiter.drawOrbit(ax, params)
    # Saturn.drawOrbit(ax, params)
    # Uranus.drawOrbit(ax, params)
    # Neptune.drawOrbit(ax, params)
    # Pluto.drawOrbit(ax, params)

    Earth.plotPointOnEVE(ax, params)

    Mercury.plotPoint(ax, params, target_date)
    Venus.plotPoint(ax, params, target_date)
    Earth.plotPoint(ax, params, target_date)
    Mars.plotPoint(ax, params, target_date)
    # Jupiter.plotPoint(ax, params, target_date)
    # Saturn.plotPoint(ax, params, target_date)
    # Uranus.plotPoint(ax, params, target_date)
    # Neptune.plotPoint(ax, params, target_date)
    # Pluto.plotPoint(ax, params, target_date)

    Mercury.textDate(ax, params, target_date)
    Venus.textDate(ax, params, target_date)
    Earth.textDate(ax, params, target_date)
    Mars.textDate(ax, params, target_date)
    # Jupiter.textDate(ax, params, target_date)
    # Saturn.textDateText(ax, params, target_date)
    # Uranus.textDateText(ax, params, target_date)
    # Neptune.textDateText(ax, params, target_date)
    # Pluto.textDateText(ax, params, target_date)

    Mercury.textAngleEVE(ax, params, target_date)
    Venus.textAngleEVE(ax, params, target_date)
    Earth.textAngleEVE(ax, params, target_date)
    Mars.textAngleEVE(ax, params, target_date)
    # Jupiter.textAngleEVE(ax, params, target_date)
    # Saturn.textAngleEVE(ax, params, target_date)
    # Uranus.textAngleEVE(ax, params, target_date)
    # Neptune.textAngleEVE(ax, params, target_date)
    # Pluto.textAngleEVE(ax, params, target_date)

    ### JAXA ###
    ## constructor : JAXA(nameStr, data_file)
    Haya2 = JAXA("Hayabusa2", "haya2_orbit_jaxa.txt")
    Ryugu = JAXA("Ryugu", "haya2_orbit_jaxa.txt")

    ## JAXA.drawOrbitJAXA(axisObj, paramsDic, begin_lp, end_lp[opt])
    Haya2.drawOrbitJAXA(ax, params, 300, 400)
    Ryugu.drawOrbitJAXA(ax, params, 300, 400)

    ## JAXA.plotPointJAXA(axisObj, paramsDic, begin_lp, end_lp[opt], days_interval[opt])
    Haya2.plotPointJAXA(ax, params, 365)
    Ryugu.plotPointJAXA(ax, params, 365)

    ## JAXA.textDateJAXA(axisObj, paramsDic, begin_lp, end_lp[opt], days_interval[opt])
    Haya2.textDateJAXA(ax, params, 365)
    Ryugu.textDateJAXA(ax, params, 365)

    ## JAXA.textAngleEVEJAXA(axisObj, paramsDic, begin_lp, end_lp[opt], days_interval[opt])
    Haya2.textAngleEVEJAXA(ax, params, 365)
    Ryugu.textAngleEVEJAXA(ax, params, 365)


    fig.suptitle("Planets and Haya2 on {0:%Y-%m-%d}".format(target_date))
    fig.savefig("planets_and_Haya2_on_{0:%Y-%m-%d}.png".format(target_date), dpi=300)

    print("end")


if __name__ == '__main__':
    main()

全体は長いので、main()と短めの"Hayabusa2","Ryugu"のインスタンス元のクラスだけを抜きだした。
これからはmain()をいじると、任意のグラフがある程度書けるようになっているかも…

久しぶりに、書いたコードをリファクタリングした。
ほぼ、書き捨てコードしか書いていないのか…

github.com