Computation¶

@schema
class Subject(dj.Manual):
    definition = """
    subject_id : varchar(16)
    ---
    species : varchar(50)
    """

@schema
class Session(dj.Manual):
    definition = """
    -> Subject
    session_idx : int32
    ---
    session_date : date
    """

    class Trial(dj.Part):
        definition = """
        -> master
        trial_idx : int32
        ---
        stimulus : varchar(50)
        response : varchar(50)
        correct : bool
        reaction_time : decimal(3,2)  # seconds
        """

@schema
class AnalysisMethod(dj.Lookup):
    definition = """
    method_name : varchar(32)
    ---
    description : varchar(255)
    """
    contents = [
        {'method_name': 'basic', 'description': 'Simple accuracy calculation'},
        {'method_name': 'weighted', 'description': 'Reaction-time weighted accuracy'},
    ]

@schema
class SessionSummary(dj.Computed):
    definition = """
    # Summary statistics for each session
    -> Session
    ---
    n_trials : int32
    n_correct : int32
    accuracy : float32
    mean_rt : float32               # mean reaction time (seconds)
    """

    def make(self, key):
        # Fetch trial data for this session (convert decimal to float)
        correct, rt = (Session.Trial & key).to_arrays('correct', 'reaction_time')
        rt = rt.astype(float)
        
        n_trials = len(correct)
        n_correct = sum(correct) if n_trials else 0
        
        # Insert computed result
        self.insert1({
            **key,
            'n_trials': n_trials,
            'n_correct': n_correct,
            'accuracy': n_correct / n_trials if n_trials else 0.0,
            'mean_rt': np.mean(rt) if n_trials else 0.0
        })

@schema
class SessionAnalysis(dj.Computed):
    definition = """
    # Analysis with configurable method
    -> Session
    -> AnalysisMethod
    ---
    score : float32
    """

    def make(self, key):
        # Fetch trial data (convert decimal to float for computation)
        correct, rt = (Session.Trial & key).to_arrays('correct', 'reaction_time')
        rt = rt.astype(float)
        
        # Apply method-specific analysis
        if key['method_name'] == 'basic':
            score = sum(correct) / len(correct) if len(correct) else 0.0
        elif key['method_name'] == 'weighted':
            # Weight correct trials by inverse reaction time
            weights = 1.0 / rt
            score = sum(correct * weights) / sum(weights) if len(correct) else 0.0
        else:
            score = 0.0
        
        self.insert1({**key, 'score': score})

@schema
class TrialAnalysis(dj.Computed):
    definition = """
    # Per-trial analysis results
    -> Session
    ---
    n_analyzed : int32
    """

    class TrialResult(dj.Part):
        definition = """
        -> master
        trial_idx : int32
        ---
        rt_percentile : float32     # reaction time percentile within session
        is_fast : bool              # below median reaction time
        """

    def make(self, key):
        # Fetch trial data
        trial_data = (Session.Trial & key).to_dicts()
        
        if not trial_data:
            self.insert1({**key, 'n_analyzed': 0})
            return
        
        # Calculate percentiles (convert decimal to float)
        rts = [float(t['reaction_time']) for t in trial_data]
        median_rt = np.median(rts)
        
        # Insert master entry
        self.insert1({**key, 'n_analyzed': len(trial_data)})
        
        # Insert part entries
        parts = []
        for t in trial_data:
            rt = float(t['reaction_time'])
            percentile = sum(r <= rt for r in rts) / len(rts) * 100
            parts.append({
                **key,
                'trial_idx': t['trial_idx'],
                'rt_percentile': float(percentile),
                'is_fast': rt < median_rt
            })
        
        self.TrialResult.insert(parts)

@schema
class SubjectSummary(dj.Computed):
    definition = """
    # Summary across all sessions for a subject
    -> Subject
    ---
    n_sessions : int32
    total_trials : int64
    overall_accuracy : float32
    """

    def make(self, key):
        # Fetch from SessionSummary (another computed table)
        summaries = (SessionSummary & key).to_dicts()
        
        n_sessions = len(summaries)
        total_trials = sum(s['n_trials'] for s in summaries)
        total_correct = sum(s['n_correct'] for s in summaries)
        
        self.insert1({
            **key,
            'n_sessions': n_sessions,
            'total_trials': total_trials,
            'overall_accuracy': total_correct / total_trials if total_trials else 0.0
        })

@schema
class QualityCheck(dj.Computed):
    definition = """
    -> Session
    ---
    passes_qc : bool
    """

    @property
    def key_source(self):
        # Only process sessions with at least 10 trials
        good_sessions = dj.U('subject_id', 'session_idx').aggr(
            Session.Trial, n='count(*)'
        ) & 'n >= 10'
        return Session & good_sessions

    def make(self, key):
        # Fetch summary stats
        summary = (SessionSummary & key).fetch1()
        
        # QC: accuracy > 50% and mean RT < 1 second
        passes = summary['accuracy'] > 0.5 and summary['mean_rt'] < 1.0
        
        self.insert1({**key, 'passes_qc': passes})

# Cleanup
schema.drop(prompt=False)