protdesign/mpnn2afserver.sh at main · scottcsh/protdesign · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
#!/bin/bash

# Usage:
# ./mpnn2afserver.sh --dir fa_dir --max_job 10 [--fa extra.fa]

while [[ $# -gt 0 ]]; do
  case $1 in
    --dir)
      DIR="$2"; shift 2 ;;
    --max_job)
      MAX_JOB="$2"; shift 2 ;;
    --fa)
      EXTRA_FA="$2"; shift 2 ;;
    *)
      echo "Usage: $0 --dir <dir> --max_job <num> [--fa extra.fa]"
      exit 1 ;;
  esac
done

[[ -z "$DIR" || -z "$MAX_JOB" ]] && { echo "--dir and --max_job required"; exit 1; }
[[ ! -d "$DIR" ]] && { echo "Invalid dir"; exit 1; }

# Optional extra chain
extra_seq=""
if [[ -n "$EXTRA_FA" ]]; then
  seq=$(grep -v '^>' "$EXTRA_FA" | tr -d '\r\n')
  extra_seq=$(cat <<EOF2
      {
        "proteinChain": {
          "sequence": "$seq",
          "count": 1
        }
      }
EOF2
)
fi

batch=1
total_job=0
json="$DIR/batch_${batch}.json"
first_in_batch=1

close_batch () {
  echo "]" >> "$json"
}

for fa in "$DIR"/*.fa; do
  [[ -e "$fa" ]] || continue
  base=$(basename "$fa" .fa)

  while IFS=$'\t' read -r tag base job_id seq; do
    ((total_job++))

    # Open a new batch file only when we actually have content to write.
    if [[ $first_in_batch -eq 1 ]]; then
      echo "[" > "$json"
    else
      echo "," >> "$json"
    fi
    first_in_batch=0

    cat >> "$json" <<EOF2
  {
    "name": "${base}_job_${job_id}",
    "modelSeeds": [],
    "sequences": [
      {
        "proteinChain": {
          "sequence": "$seq",
          "count": 1
        }
      }$( [[ -n "$extra_seq" ]] && echo ",")
$( [[ -n "$extra_seq" ]] && echo "$extra_seq" )
    ]
  }
EOF2

    if (( total_job % MAX_JOB == 0 )); then
      close_batch
      ((batch++))
      json="$DIR/batch_${batch}.json"
      first_in_batch=1
    fi

  done < <(
    awk -v base="$base" '
    /^>/ {
      if (seq != "") emit()
      seq=""
      next
    }
    {
      seq = seq $0
    }
    END {
      if (seq != "") emit()
    }
    function emit() {
      seq_count++
      if (seq_count == 1) { seq=""; return }  # skip first sequence
      file_job++
      printf "JOB\t%s\t%d\t%s\n", base, file_job, seq
      seq=""
    }
    ' "$fa"
  )
done

if [[ $first_in_batch -eq 0 ]]; then
  close_batch
fi