Charge and Polar have Ca and Cb channels inverted about timed-design HOT 2 CLOSED

encoder_list =  dataset.attrs["atom_encoder"]

for i, (pdb_code, chain_id, residue_id, _) in enumerate(data_point_batch):
    # Extract frame:
    residue_frame = np.asarray(dataset[pdb_code][chain_id][residue_id][()])

    # Check if the frame has the correct shape (final dimension is 6) for swapping:
    if residue_frame.ndim == 4 and ("Q" or "P" in encoder_list):
        # Swap only the 4th and 5th channels (index 3 and 4)
        residue_frame[..., 3], residue_frame[..., 4] = residue_frame[..., 4].copy(), residue_frame[..., 3].copy()

    X[i] = residue_frame

    # Extract residue label:
    y[i] = dataset[pdb_code][chain_id][residue_id].attrs["encoded_residue"]